时间戳

首页 > 解决方案 > pytorch dataloader default_collat​​e 参数与 to(device) 一起使用

pytorch - pytorch dataloader default_collat​​e 参数与 to(device) 一起使用

问题描述

我一直在尝试使用 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: 'list' object has no attribute 'to' 似乎默认整理的输出是长度为 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, batch_size=batch_size, shuffle=False, collate_fn=lambda x: to_device_list(x, device))

我收到以下错误： AttributeError: 'tuple' object has no attribute 'to'

请帮助如何做到这一点？

标签： pytorchbatch-processingdata-streamdataloader

解决方案

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(from ) 的每个元素传输default_collate到 cuda，最后调用 list，因为map在 python3 中被评估为惰性。

推荐阅读