PyTorch自定义Dataset场景下DataLoader返回list而非tensor问题咨询

问题原因 #

你的自定义数据集Mydataset的__getitem__方法返回了两个值：样本索引idx和图像张量img。PyTorch的DataLoader默认会将batch内所有样本的返回值按位置分别拼接，最终返回长度为2的列表：第一个元素是batch内所有样本索引拼成的张量，第二个元素是batch内所有图像拼成的张量。
你之前的遍历代码中，enumerate(test_2)返回的第一个值是batch的迭代序号，第二个值就是上述的长度为2的列表，你将其赋值给data，所以data的类型为list。

解决方法 #

有两种常用方案可选：

方案1：调整遍历的变量接收规则 #

如果你需要保留数据集返回的样本索引，直接修改遍历代码，拆分接收DataLoader返回的两个值即可：

for batch_idx, (sample_idx_list, img_tensor) in enumerate(test_2):
    # batch_idx是当前batch的迭代序号，由enumerate生成
    # sample_idx_list是batch内所有样本的索引拼成的张量
    # img_tensor就是你需要的图像batch张量
    print("batch序号：", batch_idx)
    print("图像张量：", img_tensor)
    print("图像张量类型：", type(img_tensor)) # 输出为<class 'torch.Tensor'>

如果不需要样本索引，也可以直接取列表的第二个元素：

for idx, data in enumerate(test_2):
    img_tensor = data[1]
    print(type(img_tensor))

方案2：修改数据集的返回值 #

如果你本来就不需要用到样本索引，直接修改__getitem__方法，仅返回图像即可：

def __getitem__(self, index):
     item=self.sample_list[index]
     img=Image.open(item)
     if self.transform is not None:
        img=self.transform(img)
     # 删掉返回idx的逻辑，仅返回图像张量
     return img

修改后你原来的遍历代码拿到的data就直接是张量类型。

内容的提问来源于stack exchange，提问作者Feona