Pytorch: загрузка образца изображений с помощью DataLoader

Question

Я использую стандартный DataLoader из torch.utils.data. Я создаю класс набора данных, а затем создаю DataLoader следующим образом:

train_dataset = LandmarksDataset(os.path.join(args.data, 'train'), train_transforms, split="train")
train_dataloader = data.DataLoader(train_dataset, batch_size=args.batch_size, num_workers=2,
                                   pin_memory=True, shuffle=True, drop_last=True)

Он работает отлично, но набор данных достаточно велик - 300 КБ изображений. Так что чтение изображений при использовании DataLoader занимает много времени. Так что создавать такой большой DataLoader на этапе отладки - это действительно жалко! Я просто хочу проверить свою гипотезу и сделать это быстро! Для этого мне не нужно загружать весь набор данных.

Я пытаюсь найти способ Как загрузить только небольшую фиксированную часть набора данных без создания dataLoader для всего набора данных? В настоящий момент все мои идеи просто создают другую папку, копируют некоторая часть изображений здесь и использовать на ней конвейер. Но я полагаю, что Pytorch достаточно умен, чтобы иметь несколько встроенных методов для загрузки только части изображений из большого набора данных. Вы можете мне посоветовать, как это сделать?

Answer 1

Насколько мне известно, не существует механизма, который делал бы это за вас. Ваша проблема заключается в классе LandmarksDataset в том месте, где вы читаете пути к папке с данными вашего поезда. Я предполагаю os.listdir(train_data_folder).

Вместо этого вы можете использовать более эффективный способ. os.scandir(train_data_folder) возвращает генератор, а вызов next() на нем даст вам пути к вашим изображениям в данных поезда. Таким образом, вы можете вызывать next () сколько угодно раз, не изменяя структуру папки с данными поезда, и создавать ее подмножество.