@matweykai
Standart coder ; )

Как эффективно загружать картинки при обучении в PyTorch?

Решаю соревнование на Kaggle, раньше я не работал с картинками, кроме заданий на курсах. Решаю задачу классификации изображений, их примерно 9 ГБ (415к файлов). Моя модель очень долго обучается, и я заметил, что основное время уходит на загрузку данных. Сейчас я делаю это так:

class MushroomDataset(Dataset):
    def __init__(self, mapping_ds: pd.DataFrame, pict_folder: str, transforms=None):
        self.mapping_ds = mapping_ds
        self.pict_folder = pict_folder
        self.transforms = transforms
        
    def __getitem__(self, ind: int):        
        cur_obj = self.mapping_ds.iloc[ind]
        
        pict_path = path.join(self.pict_folder, cur_obj.image_path[:-4] + '.jpg')
        
        img = Image.open(pict_path)
        
        if self.transforms:
            img = self.transforms(img)
            
        return img, cur_obj.class_id
    
    def __len__(self):
        return self.mapping_ds.shape[0]


Модель проходит одну эпоху за час, что мне кажется ненормальным. Если интересно, то код для обучения такой:

def train_model(model, optimizer, loss, train_loader, val_loader, scheduler=None, epoch_num=10):
    loss_history = list()
    train_hist = list()
    val_hist = list()
    
    model.to(device)
    
    for epoch_ind in range(epoch_num):
        model.train()
        
        corr_predicted_obj = 0
        loss_accum = 0
        # Training
        for i_step, (x, y) in enumerate(train_loader, 1):
            x = x.to(device)
            y = y.to(device)
            
            pred = model(x)
            
            loss_val = loss(pred, y)
            
            optimizer.zero_grad()
            loss_val.backward()
            optimizer.step()
            
            loss_accum += loss_val
            
            pred_labels = torch.argmax(pred, dim=1)
            corr_predicted_obj += torch.sum(pred_labels == y)
            
            
        loss_history.append(loss_accum / i_step)
        train_hist.append(corr_predicted_obj / len(train_loader.dataset))
        
        # Validation
        val_corr_predicted_obj = 0
        
        model.eval()
        
        for x, y in val_loader:
            x = x.to(device)
            y = y.to(device)
            
            pred = model(x)
            
            pred_labels = torch.argmax(pred, dim=1)

            val_corr_predicted_obj += torch.sum(pred_labels == y)
        
        val_hist.append(val_corr_predicted_obj / len(val_loader.dataset))
        
        # Printing model progress
        print(f"Epoch -> {epoch_ind + 1}  Loss value: {round(loss_history[-1].item(), 5)}  " + \
              f"Train score: {round(train_hist[-1].item(), 4) * 100} %  " + \
              f"Validation score: {round(val_hist[-1].item(), 4) * 100}")
        
        # Scheduler step
        if scheduler:
            scheduler.step()
        
    return loss_history, train_hist, val_hist

А модель вот:

frez_model = models.densenet121(pretrained=True)

for layer in frez_model.parameters():
    layer.requires_grad = False

frez_model.classifier = nn.Linear(frez_model.classifier.in_features, len(train_ds.class_id.unique()))

lr = 1e-2

optimizer = optim.SGD(frez_model.parameters(), lr=lr)
scheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.5)

loss = nn.CrossEntropyLoss()

loss_history, train_hist, val_hist = train_model(frez_model, optimizer, loss, train_loader, valid_loader, scheduler=scheduler)


Перерыл уже множество Kaggle Kernel на этом и похожем соревновании(тоже классификация картинок и возникает та же проблема). Как обучать всё это дело быстрее? (Если что обучение происходит также на железе Kaggle и на их GPU)
  • Вопрос задан
  • 164 просмотра
Пригласить эксперта
Ваш ответ на вопрос

Войдите, чтобы написать ответ

Похожие вопросы