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

Question

如果我们使用

Dataset

和

Dataloader

类的组合（如下所示），我必须使用 .to() 或

.cuda()

将数据显式加载到

GPU

上。有没有办法指示数据加载器自动/隐式地执行此操作？

理解/重现场景的代码：

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

一个可能的解决方案位于 这个 PyTorch GitHub 存储库。问题（在发布此问题时仍处于开放状态），但是，当数据加载器必须返回多个数据项时，我无法使其工作！

Answer 1

您可以修改

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

会更有效。您尝试过这个选项吗？
有关更多信息，请参阅内存固定。

---

更新（2021 年 2 月 8 日）
这篇文章让我了解了训练期间花费的“数据到模型”时间。 我比较了三种选择：

1. DataLoader

   适用于 CPU，仅在检索批次后数据才会移至 GPU。
2. 与 (1) 相同，但

   pin_memory=True

   中带有

   DataLoader

   。
3. 建议使用

   collate_fn

   将数据移动到 GPU 的方法。

从我有限的实验来看，第二个选项似乎表现最好（但幅度不是很大）。
第三个选项需要对数据加载进程的

start_method

大惊小怪，并且它似乎会在每个纪元开始时产生开销。

Answer 2

初始化DataSet时可以先将数据传递给GPU。

参见https://stackoverflow.com/a/68972606/8314159