PyTorch: при загрузке пакетных данных с помощью Dataloader, как автоматически передавать данные в графический процессор

Если мы используем комбинацию классов Dataset и Dataloader (как показано ниже), я должен явно загрузить данные в GPU, используя .to() или .cuda(). Есть ли способ указать загрузчику данных делать это автоматически / неявно?

Код для понимания / воспроизведения сценария:

from torch.utils.data import Dataset, DataLoader
import numpy as np

class DemoData(Dataset):
    def __init__(self, limit):
        super(DemoData, self).__init__()
        self.data = np.arange(limit)

    def __len__(self):
        return self.data.shape[0]

    def __getitem__(self, idx):
        return (self.data[idx], self.data[idx]*100)

demo = DemoData(100)

loader = DataLoader(demo, batch_size=50, shuffle=True)

for i, (i1, i2) in enumerate(loader):
    print('Batch Index: {}'.format(i))
    print('Shape of data item 1: {}; shape of data item 2: {}'.format(i1.shape, i2.shape))
    # i1, i2 = i1.to('cuda:0'), i2.to('cuda:0')
    print('Device of data item 1: {}; device of data item 2: {}\n'.format(i1.device, i2.device))

Будет выведено следующее: примечание - без явной инструкции передачи устройства данные загружаются в ЦП:

Batch Index: 0
Shape of data item 1: torch.Size([50]); shape of data item 2: torch.Size([50])
Device of data item 1: cpu; device of data item 2: cpu

Batch Index: 1
Shape of data item 1: torch.Size([50]); shape of data item 2: torch.Size([50])
Device of data item 1: cpu; device of data item 2: cpu

Возможное решение находится в репозитории PyTorch на GitHub. Проблема (все еще открыта на момент публикации этого вопроса), но я не могу заставить ее работать, когда загрузчик данных должен возвращать несколько элементов данных!


python-3.x pytorch gpu pytorch-dataloader
person anurag    schedule 28.01.2021    source источник
comment
Никто не сталкивался с этой проблемой?   -  person anurag    schedule 01.02.2021
comment
Параметр collate-fn бесполезен, если набор данных возвращает кортеж значений в каждом пакете. Так что все еще ищем лучшие обходные пути!   -  person anurag    schedule 02.02.2021

Ответы (1)


arrow_upward
2
arrow_downward

Вы можете изменить collate_fn для обработки нескольких элементов одновременно:

from torch.utils.data.dataloader import default_collate

device = torch.device('cuda:0')  # or whatever device/cpu you like

# the new collate function is quite generic
loader = DataLoader(demo, batch_size=50, shuffle=True, 
                    collate_fn=lambda x: tuple(x_.to(device) for x_ in default_collate(x)))

Обратите внимание: если вы хотите иметь несколько рабочих для загрузчика данных, вам нужно добавить

torch.multiprocessing.set_start_method('spawn')

после вашего if __name__ == '__main__' (см. эту проблему).

Сказав это, похоже, что использование pin_memory=True в вашем DataLoader было бы намного эффективнее. Вы пробовали этот вариант?
См. закрепление памяти для больше информации.

Обновление (8 февраля 2021 г.)
Эта статья заставила меня взглянуть на время, затраченное на преобразование данных в модель во время обучения. Я сравнил три альтернативы:

  1. DataLoader работает на CPU и только после получения пакета данные перемещаются в GPU.
  2. То же, что (1), но с pin_memory=True в DataLoader.
  3. Предлагаемый метод использования collate_fn для переноса данных на GPU.

Из моих ограниченных экспериментов кажется, что второй вариант работает лучше всего (но не с большим отрывом).
Третий вариант потребовал возиться с start_method процессами загрузчика данных, и, похоже, в начале каждую эпоху.

person Shai    schedule 02.02.2021
comment
spawn инструкция у меня не сработала, она специфична для конкретной версии? - person anurag; 08.02.2021
comment
@anurag Я не знаю. - person Shai; 08.02.2021