PyTorch自定义数据集__getitem__解包错误及图片数组维度问题求助

报错根因 #

你的问题核心是数据集里的样本格式不统一：
一部分样本是符合预期的(图片数组, 标签)二元组，所以前几次调用__getitem__可以正常跑，等DataLoader加载到四元组格式的异常样本时，就会触发解包错误。

修复代码 #

先导入numpy，修改__getitem__方法兼容两种格式即可：

import numpy as np

def __getitem__(self, index):
    data_item = self.data[index]
    # 适配两种样本格式
    if len(data_item) == 2:
        img, label = data_item
    else:
        part1, part2, part3, label = data_item
        # 把part3的列表转为numpy数组，形状为(150, 312)
        part3_arr = np.array(part3)
        # 以下合并逻辑可以根据你实际需要的图片维度调整，示例给两种常用方式：
        # 方案1：直接堆叠三个部分得到3维数组，适配多通道输入要求
        img = np.stack([
            part1, 
            np.array(part2), 
            part3_arr.flatten()
        ], axis=0)
        # 方案2：如果只需要用part1作为输入，直接reshape到2维即可
        # img = part1.reshape(32, 64) # 2048=32*64，按需调整尺寸
        
        # 如果需要适配PIL的(H,W,C)维度顺序，加一行转置即可
        # img = np.transpose(img, (1,2,0))
    return img, int(label)

后续优化建议 #

这个兼容方案只是临时解决运行报错，你最好回头排查下数据集的生成逻辑：

• 确认为什么会同时生成两种格式的样本
• 从数据生产侧统一格式，比下游加载时做兼容更稳妥
• 可以在生成数据集时加格式校验，过滤掉异常样本，避免运行时中途崩溃

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