PyTorch:使用Dataloader加载批量数据时,如何自动将数据传输到GPU

PyTorch: while loading batched data using Dataloader, how to transfer the data to GPU automatically

如果我们使用 DatasetDataloader 类 的组合(如下所示),我必须明确地将数据加载到 GPU 使用 .to().cuda()。有没有办法指示数据加载器执行此操作automatically/implicitly?

understand/reproduce场景代码:

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))

这将输出以下内容;注意 - 没有明确的设备传输指令,数据被加载到 CPU:

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

可能的解决方案在 this PyTorch GitHub repo. Issue发布此问题时仍处于打开状态),但是,当数据加载器必须 return 多个数据项!

您可以修改 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__' 之后(参见 this issue)。

话虽如此,在 DataLoader 中使用 pin_memory=True 似乎会更有效率。您尝试过这个选项吗?
有关详细信息,请参阅 memory pinning


更新(2021 年 2 月 8 日)
post 让我查看了我在训练期间花费的“数据到模型”时间。 我比较了三种选择:

  1. DataLoader 在 CPU 上工作,只有在检索到批次后,数据才会移动到 GPU。
  2. 与 (1) 相同,但 DataLoader 中有 pin_memory=True
  3. 建议使用collate_fn将数据移动到GPU的方法。

从我有限的实验来看,第二个选项似乎表现最好(但幅度不大)。
第三个选项需要对数据加载进程的 start_method 大惊小怪,而且它似乎会在每个 epoch 开始时产生开销。