PyTorch:在使用Dataloader加载批处理数据时，如何将数据自动传输到GPU
PyTorch:在使用Dataloader加载批处理数据时，如何将数据自动传输到GPU
提问于 2021-01-28 06:56:53
如果我们使用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回购。问题(在这个问题发布时仍然是开放的)，但是，当数据处理程序必须返回多个数据项时，我无法使它工作！
回答 1
Stack Overflow用户
回答已采纳
发布于 2021-02-02 15:17:52
您可以修改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日)
这篇文章让我审视了我在训练中所花的“数据到模型”的时间。我比较了三种选择：
DataLoader在CPU上工作，只有在获取批处理数据后才能将数据移动到GPU。- 与(1)相同，但与
pin_memory=True在DataLoader中相同。 - 提出了利用
collate_fn将数据移动到GPU的方法。
从我有限的实验来看，第二种选择似乎表现最好(但不是以很大的优势)。
第三种选择需要对数据加载器进程的start_method进行复杂的处理，并且在每个时代开始时似乎都会产生开销。
页面原文内容由Stack Overflow提供。腾讯云小微IT领域专用引擎提供翻译支持
原文链接：
https://stackoverflow.com/questions/65932328
复制相关文章
相似问题
腾讯云开发者
