PyTorch自定义数据集类教程

本文深入剖析了PyTorch自定义Dataset类的五大核心陷阱：__getitem__必须返回张量而非NumPy数组（否则torch.stack()会直接报错）、__len__必须严格匹配真实可索引样本数（否则训练到末尾才崩溃且难以调试）、数据增强务必在__getitem__中每次动态生成随机结果（避免状态复用导致全批样本增强一致）、多进程加载时异常被静默吞没的排查技巧（临时设num_workers=0或加全局异常捕获），以及极易被忽视的文件路径、解码失败、标签错位等边界情况防御——每一条都直击实战中最让人抓狂的“看似正常却突然崩盘”的痛点，帮你避开那些让模型跑三小时后才报错的隐形坑。

为什么继承 Dataset 后 __getitem__ 必须返回张量而不是 NumPy 数组

PyTorch 的 DataLoader 默认会调用 torch.stack() 合并一个 batch 的样本，而 torch.stack() 要求所有输入必须是同类型的 torch.Tensor。如果 __getitem__ 返回的是 np.ndarray 或 Python 原生类型，DataLoader 在 collate 阶段大概率报错：TypeError: expected Tensor as element 0 in argument 0, but got numpy.ndarray。

• 实操建议：在 __getitem__ 内部就完成 torch.tensor(img) 或 torch.from_numpy(arr) 转换，别留到 collate_fn 里补
• 注意 torch.from_numpy() 是零拷贝，但要求数组内存连续；不放心就加 .copy() 或用 torch.tensor()
• 图像常用库（PIL / OpenCV）读出的数据默认不是 float32，直接转 tensor 后可能值域不对（比如 0–255 整数变成长整型），后续 Normalize 会失效

__len__ 返回值不等于实际可索引样本数时会发生什么

__len__ 不是“声明长度”，而是 DataLoader 构建索引范围的唯一依据。如果它返回 100，但 __getitem__(99) 报 IndexError 或返回脏数据，训练会在第 100 个 batch 突然崩掉，错误堆栈往往只显示在 DataLoader worker 内部，难以定位。

• 常见错误：用文件列表初始化时忘了过滤掉损坏图片，len 按原始文件数算，但 __getitem__ 加载失败后没抛异常也没降级处理
• 安全做法：在 __init__ 中预扫描路径，只保留能成功打开/解析的样本路径，再设 self.samples = valid_paths
• 调试技巧：临时把 __len__ 改成 min(10, len(self.samples)) 快速验证前几个样本是否真能走通

自定义 __getitem__ 时要不要做数据增强？什么时候做

增强逻辑放在 __getitem__ 里是标准做法，但关键在于「增强必须每次调用都产生新结果」。如果用了固定随机种子、或复用了某个 transform 实例却没重置状态，会导致所有样本增强效果一致（比如全部水平翻转，或全部裁剪同一区域）。

• 别在 __init__ 里写 self.transform = transforms.RandomHorizontalFlip(p=1.0) —— 这个对象内部有状态，多次调用不会重新采样
• 正确方式：用 torchvision.transforms.Compose 包裹原子操作，或者自己写函数，在 __getitem__ 内部调用时生成新随机数（如 if random.random() > 0.5: img = img.transpose(Image.FLIP_LEFT_RIGHT)）
• 性能提示：CPU 上做 heavy augment（如大图 resize + 弹性形变）会拖慢 DataLoader，考虑用 num_workers > 0 分摊，但要小心多进程下随机种子同步问题

继承 Dataset 后，__getitem__ 报错却看不到完整 traceback

当 DataLoader 开了 num_workers > 0，worker 进程中发生的异常默认被吞掉，主进程只收到模糊提示如 BrokenPipeError 或直接卡死。真实错误藏在 worker 子进程中，不显式捕获根本看不到。

• 快速诊断：临时把 DataLoader 的 num_workers=0，让一切在主线程执行，此时 __getitem__ 的异常会原样抛出
• 长期方案：在 __getitem__ 最外层加 try...except Exception as e: print(f"Sample {idx} failed: {e}"); raise，避免静默失败
• 注意：某些异常（如 PIL 读取损坏 JPEG）可能不继承 Exception，用裸 except: 更保险（仅调试期）

实际写的时候最容易漏的是：__getitem__ 里没处理路径不存在、图像解码失败、标签文件格式错位这些边界情况，等模型训到第 3 个 epoch 才崩，回溯成本远高于初始化时多写三行防御代码。

终于介绍完啦！小伙伴们，这篇关于《PyTorch自定义数据集类教程》的介绍应该让你收获多多了吧！欢迎大家收藏或分享给更多需要学习的朋友吧~golang学习网公众号也会发布文章相关知识，快来关注吧！