PyTorch自定义数据集加载触发DataLoader worker KeyError问题咨询

报错根因与修复方案 #

你当前遇到的KeyError和图像尺寸没有任何关系，错误栈明确指向pandas索引取值失败：train_test_split拆分数据集后，X_train的索引保留了原train_data的全局索引，不是从0开始的连续整数序列。而Dataset的__getitem__传入的idx是0到数据集长度-1的连续下标，用这个下标去取索引不连续的image_ids Series时，就会找不到对应索引触发报错。

修复方法很简单，初始化Dataset的时候把id和标签转成列表或者重置索引即可：

# 修改CustomImageDataset的__init__方法
def __init__(self, dataset, transforms=None, target_transforms=None):
    # 转成列表后就可以用下标正常取值
    self.image_ids = dataset.image_id.tolist()
    self.image_labels = dataset.label.tolist()
    self.img_dir = 'Data/images'
    self.transforms = transforms
    self.target_transforms = target_transforms

问题1：异尺寸图像的训练前预处理方案 #

针对你提到的(512,512,3)三通道图和(1024,1024)单通道图混合的情况，常用预处理方案有以下几种：

• 尺寸统一处理：
  可以用torchvision.transforms.Resize将所有图像缩放到同一尺寸，比如统一缩到512512；如果不想损失图像内容，可以结合Pad填充到最大尺寸10241024，填充值可以选择0或者数据集的像素均值。训练阶段也可以用RandomResizedCrop做数据增强，同时保证输出尺寸一致。
• 通道数统一处理：
  新增通道转换逻辑，把单通道图转成三通道，可直接使用transforms.Grayscale(num_output_channels=3)实现，避免后续拼batch时通道数不匹配报错。
• 动态尺寸适配：
  如果模型结构中加入了自适应池化层（比如AdaptiveAvgPool2d），不需要强制所有样本尺寸完全相同，只要保证同一个batch内的图像尺寸一致即可，适合目标检测、分割等对原图尺寸敏感的任务。

问题2：该类KeyError的其他触发原因 #

排除当前索引不连续的问题后，DataLoader worker进程抛出的KeyError还有以下常见诱因：

• 标注文件中的图像id与实际存储的文件名不匹配，部分id找不到对应的图像文件
• 自定义Dataset中使用字典存储标注信息，部分样本对应的键不存在
• 多进程加载时，主进程和子进程的数据集索引没有对齐，比如在Dataset的__init__方法中做了随机打乱，子进程复制数据集时产生了不同的打乱结果
• 数据集存在脏数据，比如部分样本的id为空、标签缺失，加载时无法找到对应键值

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