Как правильно реализовать t.utils.data.Dataloader

Question

Я работаю над проектом, который требует, чтобы я получил данные, метку из файла csv, прежде чем идти на обучение в PyTorch. Структура CSV-файла:

изображение, животное, вымерший

blabla.png, 0, 1

..............

и так далее..

Я создал класс для получения набора данных, например:

      from torch.utils.data import Dataset

class CallDataset(Dataset):
    # TODO implement the Dataset class according to the description

    def __init__(self, data, mode):
        #self.img, self.lbl = data
        self.data = data
        self.mode = mode
        self._transform = transforms.Compose([
            transforms.ToPILImage(),
            transforms.ToTensor(),
            transforms.Normalize(mean=train_mean, std=train_std)
        ])
        # self.tab = pd.read_csv(self.data, sep=';')
        self.image = data.iloc[0:, 0]
        self.label = data.iloc[0:, 1:]
        #self.image = self.img
        #self.label = self.lbl


    def __len__(self):
        return len(self.image)

    def __getitem__(self, index):
        if self.mode == 'train':
            x = imread(self.image[index])
            sample = gray2rgb(x)
            transformed_img = self._transform(sample)

            transformed_label = self.label.to_numpy()
            transformed_label = transformed_label[index]
            transformed_label = torch.from_numpy(transformed_label)

            return (transformed_img, transformed_label)
        
        if self.mode == 'val':
            x = imread(self.image[index])
            sample = gray2rgb(x)
            transformed_img = self._transform(sample)

            transformed_label = self.label.to_numpy()
            transformed_label = transformed_label[index]
            transformed_label = torch.from_numpy(transformed_label)

            return (transformed_img, transformed_label)

Итак, идея состоит в том, чтобы вернуть кортеж из двух тензоров torch. Первый - это тензор преобразованного изображения RGB формы (1, 3, высота, ширина), а второй - тензор факела, состоящий из меток (независимо от того, является ли это животным и вымершим). Его форма, если (1, 2).

Файл csv и соответствующая папка изображений содержат 1000 изображений и 1000 соответствующих наборов этикеток.

Для тренировок пытаюсь называть их так:

         def train_epoch(self):
        
        self._model = self._model.train()
        l = 0
        c = 0
        for x, y_true in self._train_dl:
            c +=1
            if self._cuda:
                
                x = x.to('cuda')
                y_true = y_true.to('cuda')
            l += self.train_step(x, y_true)
        avg_loss = l/c
        return avg_loss

Self.train_dl определяется следующим образом:

      def train_dataset():
    train_data = CallDataset(train, 'train')
    return train_data
train_dl = t.utils.data.DataLoader(train_dataset(), batch_size=200, shuffle=True)

Итак, идея состоит в том, чтобы получать на каждой итерации изображение и соответствующую метку, отправлять их в cuda, обучать их (train_step - это функция, в которой я вызвал модель, побежал вперед, получил убытки, побежал назад и т. Начальная потеря l = 0 - это то место, где я определил l. c в основном хранит количество взаимодействий, которое мне нужно, чтобы найти среднюю потерю (l,c меня не беспокоит на данный момент).

У меня вопрос, правильно ли я использую Dataloader. Кроме того, дайте мне знать, как лучше всего их использовать с пояснениями.