pytorch dataloader default_collat​​e参数与todevice一起使用

编程问答 2022-06-15

问题描述

我一直在尝试使用to（device）将to（device）集成到我的数据加载器中，如https://github.com/pytorch/pytorch/issues/11372

所示

我通过以下方式在FashionMNIST上定义了它：

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
batch_size = 32
trainset = datasets.FashionMNIST('~/.pytorch/F_MNIST_data/',download=True,train=True,transform=transforms.ToTensor())
rain_loader = torch.utils.data.DataLoader(trainset,batch_size=batch_size,shuffle=False,collate_fn=lambda x: default_collate(x).to(device))

但是我得到以下错误： AttributeError：“列表”对象没有属性“至” 似乎默认排序规则的输出是长度为2的列表，第一个元素是图像张量，第二个元素是标签张量（因为它的collat​​e_fn = None输出next（iter（train_loader））），所以我尝试了具有以下已定义的功能：

def to_device_list(l,device):
    return [l[0].to(device),l[1].to(device)]
train_loader = torch.utils.data.DataLoader(trainset,collate_fn=lambda x: to_device_list(x,device))

然后出现以下错误： AttributeError：“元组”对象没有属性“至”

有什么帮助吗？

解决方法

fashion mnist 数据集返回 img 和 target 的元组，其中 img 是张量，目标是 int 类值。

现在，您的 dataloader 从 dataset 类中获取批量大小的样本以获取样本列表。请注意，此示例列表现在为 List[Tuple[Tensor,int]]（此处使用输入注释）。然后它调用 collat​​e 函数将 List[Tuple[Tensor,int]] 转换为 List[Tensor]，其中这个列表有 2 个张量。第一个张量是大小为 [32,1,28,28] 的图像的堆叠数组，其中 32 是批量大小，第二个张量是 int 值（类标签）的张量数组。

default_collate 函数，只是将结构数组转换为数组结构。

现在，当您使用 collate_fn=lambda x: default_collate(x).to(device) 时，请注意 default_collat​​e 会返回一个张量列表。所以在列表上调用 .to 不起作用，应该在列表的所有元素上调用。

解决方案 使用

collate_fn=lambda x: list(map(lambda x: x.to(device),default_collate(x))))

map 函数将 list 的每个元素（从 default_collate）转移到 cuda，最后调用 list，因为 map 在 python3 中被评估为惰性。

batch-processingdata-streampytorch