1. 从零手写自定义数据集：为什么这件事值得认真对待
搞深度学习的人迟早会撞上一个坎：手头的数据不是现成的 MNIST、CIFAR 或者 ImageNet 那种规规矩矩的文件夹结构，而是一堆命名混乱的图片、几个格式各异的标注文件、甚至是从数据库里导出来的 CSV。这时候torch.utils.data.Dataset就成了绕不开的东西。我见过太多项目在原型阶段用ImageFolder跑得飞快，一到真实数据就卡在数据加载这一步，最后花两天时间重写数据管道，比搭模型的时间还长。
pytorch-custom dataset这个主题说白了就是：如何把你自己那份"不听话"的数据，包装成 PyTorch 训练循环能直接消费的标准接口。它解决的核心问题是数据与模型之间的适配层问题——模型只认张量和批次，而你的原始数据可能是任意形态。适合谁来参考？我认为三类人最需要：一是刚入门、只会用现成数据集跑 demo 的新手；二是手里有业务数据、需要快速搭训练管道的算法工程师；三是做多模态或者特殊任务、标准数据集根本覆盖不了场景的研究者。
这篇文章我不打算复述官方文档里那些__init__、__len__、__getitem__的基础定义，那些东西看一遍就懂。我要讲的是实际写自定义数据集时真正会踩的坑：为什么你的 DataLoader 越跑越慢、为什么多进程下随机种子失效、为什么 transform 写错位置会导致数据增强在验证集上也生效、为什么内存会莫名其妙爆掉。这些经验在文档里基本找不到，但每一个都能让你多调试半天。
先给一个整体认知：一个自定义数据集的质量，决定了你整个训练流程的上限。模型再强，如果数据管道有瓶颈，GPU 利用率上不去，训练时间翻倍；如果数据增强逻辑有 bug，模型学到的分布就是错的，指标虚高但上线就崩。所以这个主题不是"会用就行"的入门知识，而是值得反复打磨的工程基本功。
2. 整体设计思路：自定义数据集到底在解决什么问题
2.1 数据管道的三层结构拆解
PyTorch 的数据加载体系其实是三层：最底层是Dataset，负责"单条样本怎么取"；中间是Sampler，负责"按什么顺序取"；最上层是DataLoader，负责"怎么批量取、用几个进程取"。很多人写自定义数据集时只盯着第一层，结果后面两层出问题时完全不知道从哪查。
我习惯把这个结构类比成餐厅后厨：Dataset是配菜员，负责把一道菜的原料备齐；Sampler是传菜顺序，决定先做哪桌；DataLoader是出餐节奏，控制一次出几份、几个灶台同时开工。配菜员再快，如果出餐节奏乱了，整个后厨照样堵。所以设计自定义数据集时，脑子里要有这个全局图，而不是孤立地写一个类。
具体到代码层面，Dataset只需要实现两个方法：__len__返回样本总数，__getitem__接收索引返回一条样本。听起来简单，但魔鬼全在细节里。__getitem__的返回值结构直接决定了后面collate_fn怎么写、模型输入怎么组织；__len__如果返回错了，Sampler的索引范围就错了，轻则报 IndexError，重则静默丢样本。
2.2 为什么不用现成的 ImageFolder 或 TensorDataset
有人会问：ImageFolder不是挺好用吗，为什么还要自己写？关键在于ImageFolder假设了"一个类别一个文件夹"的硬性结构，而且标签是从文件夹名推断的。真实场景里，你的标注可能是 CSV 里的一列、可能是 JSON 里的嵌套字段、可能一张图对应多个标签、可能标签本身是连续值而非离散类别。这些情况ImageFolder全都搞不定。
TensorDataset则是另一个极端：它要求你事先把所有数据转成张量放进内存。小数据集没问题，但图像、视频、长文本这类数据，全量加载进内存基本等于自杀。自定义数据集的最大价值就是惰性加载——只在__getitem__被调用时才去读那一条数据，内存占用只跟 batch size 相关，跟数据集总大小无关。
还有一个容易被忽略的点：自定义数据集让你能精确控制数据增强的时机和范围。用现成封装时，增强逻辑往往藏在内部，你想针对某个类别做特殊处理、想在验证时关掉某些增强，都很别扭。自己写的话，每一行逻辑都在你眼皮底下，调试和优化都有抓手。
2.3 惰性加载与预加载的取舍逻辑
惰性加载不是银弹。如果你的数据存在机械硬盘上、单条读取耗时几十毫秒，那每个 batch 都要等 IO，GPU 大量时间在空转。这时候就要考虑预加载或者缓存策略。我的经验是分三档处理：
- 数据总量小于内存的 30%，直接全量预加载到内存，用列表或字典存，
__getitem__就是一次索引，快得飞起。 - 数据量中等、但单条读取慢（比如需要解码的大图），用 LRU 缓存最近访问的若干条，兼顾内存和速度。
- 数据量巨大、无法全量加载，老老实实惰性读取，但要把 IO 并行化，靠
DataLoader的多进程来掩盖延迟。
这个取舍没有标准答案，取决于你的硬件和数据特征。我一般会先写惰性版本跑通逻辑，再用 profiler 看数据加载占总时间的比例，超过 20% 就考虑加缓存。盲目预加载导致 OOM 的教训我吃过不止一次，所以现在默认从惰性开始。
3. 核心细节解析：写对getitem的几个关键决策
3.1 返回值结构设计与 collate_fn 的配合
__getitem__返回什么，是整个数据管道里最重要的设计决策。最常见的返回是(data, label)元组，但真实任务里往往更复杂：目标检测要返回图像、边界框列表、类别列表；分割任务要返回图像和掩码；多任务学习可能要返回一个字典。返回值结构一旦定下来，collate_fn就得跟着适配。
默认的collate_fn只能处理"每条样本结构完全一致、且能直接 stack"的情况。如果你的样本里边界框数量不一致，默认 collate 会直接报错，因为不同长度的列表没法 stack。这时候必须自定义collate_fn，手动把变长部分整理成列表或者 padding 成定长张量。
我踩过的一个坑是：早期图省事，__getitem__返回字典，键名用了中文或者带空格的字符串，结果多进程下序列化出问题。后来统一规范：键名只用小写英文加下划线，值只用张量、数字、字符串、列表这几种基础类型，避免嵌套过深。这个约定看起来死板，但省了无数调试时间。
3.2 索引映射：从全局索引到真实数据定位
__getitem__收到的index是 0 到len-1的连续整数，但你的真实数据可能不是这样组织的。比如数据存在一个字典里、键是文件名的哈希；或者数据分片存在多个文件里，需要先算出属于哪个分片、再算分片内偏移。这个映射逻辑必须写对，而且要写得高效。
我的做法是在__init__里预先构建一个索引列表，把"全局索引"映射到"数据定位信息"。比如一个列表，第 i 个元素是第 i 条样本的文件路径或者 (分片号, 偏移) 元组。这样__getitem__里就是一次列表索引加一次读取，逻辑清晰且 O(1)。千万不要在__getitem__里做线性搜索或者复杂计算，那个方法会被调用几万几十万次，任何多余操作都会被放大。
还有一个细节：如果数据集需要过滤（比如跳过损坏文件、跳过某些类别），过滤逻辑要放在__init__构建索引列表时做，而不是在__getitem__里判断后返回 None。返回 None 会让collate_fn崩溃，而且过滤后索引和长度对不上，非常容易出隐蔽 bug。
3.3 transform 的位置：放在 Dataset 里还是外面
这是一个争论很多的问题。__getitem__里做 transform，还是把原始数据返回、在训练循环里做？我的观点很明确：数据增强放在 Dataset 内部，但要做成可配置的。
原因有三。第一，多进程加载时，transform 在 worker 进程里执行，能真正并行，如果放到主进程的训练循环里，就变成串行了，白白浪费多核。第二，Dataset 内部做增强，逻辑内聚，换任务时只改一个类，不用动训练代码。第三，验证集和训练集用不同的 transform 实例，天然隔离，不会出现"验证集也做了随机裁剪"这种低级错误。
但要注意，transform 里如果有随机性，必须保证每个 worker 的随机种子不同，否则所有 worker 会产生相同的随机序列，增强效果大打折扣。PyTorch 的 DataLoader 在 worker 初始化时会设置种子，但如果你在 Dataset 里自己用了random或numpy.random，需要手动在worker_init_fn里重新播种。这个坑我在多进程训练时踩过，表现为"数据增强好像没生效"，查了半天才发现是种子问题。
4. 实操过程：一个完整可复现的自定义数据集实现
4.1 场景设定与数据组织
假设我们有一个图像分类任务，数据组织方式是：所有图片放在一个目录下，标注信息在一个 CSV 文件里，包含filename和label两列。这个结构在真实项目里极其常见，比"一类别一文件夹"更贴近实际。我们还要支持可选的 transform、可选的缓存、以及训练/验证模式的区分。
先明确几个设计约束：图片路径和标签在__init__时一次性读入并校验；__getitem__负责读图、转 RGB、应用 transform、返回(tensor, label)；损坏文件在初始化时就被过滤掉并记录日志；支持传入transform=None时返回 PIL 图像（方便调试）。
4.2 初始化阶段的校验与索引构建
初始化阶段做的事情比想象中多。读 CSV 只是第一步，还要做几件事：检查文件是否存在、检查标签是否合法、构建索引列表、记录统计信息。我习惯在初始化时打印一份摘要：总样本数、过滤掉的损坏样本数、类别分布。这份摘要能帮你快速发现数据问题，比如某个类别样本极少、或者路径配置错了导致全部文件找不到。
import os import csv import logging from PIL import Image from torch.utils.data import Dataset class CsvImageDataset(Dataset): def __init__(self, csv_path, image_dir, transform=None, label_to_idx=None): self.image_dir = image_dir self.transform = transform self.samples = [] self.label_to_idx = label_to_idx or {} skipped = 0 with open(csv_path, newline='', encoding='utf-8') as f: reader = csv.DictReader(f) for row in reader: fname = row['filename'].strip() label = row['label'].strip() full_path = os.path.join(image_dir, fname) if not os.path.isfile(full_path): skipped += 1 continue if label not in self.label_to_idx: self.label_to_idx[label] = len(self.label_to_idx) self.samples.append((full_path, self.label_to_idx[label])) logging.info(f"loaded {len(self.samples)} samples, skipped {skipped}, " f"{len(self.label_to_idx)} classes")这段代码有几个刻意的设计。label_to_idx允许外部传入，这样训练集和验证集能共享同一套标签映射，不会出现"训练集里猫是 0、验证集里猫是 1"的灾难。跳过不存在文件时只计数不报错，因为真实数据里总有几条脏数据，为了它们中断整个流程不值得，但必须记录数量让你心里有数。
4.3getitem的健壮实现与异常兜底
__getitem__是热路径，要快，但也不能裸奔。读图可能失败（文件损坏、格式不支持），transform 可能抛异常（尺寸不对、通道数不对）。我的策略是：读图失败时返回一个占位样本并记录，而不是直接抛异常中断训练。占位样本用全零张量加一个特殊标记，训练时可以选择忽略。
def __len__(self): return len(self.samples) def __getitem__(self, index): path, label = self.samples[index] try: img = Image.open(path).convert('RGB') except Exception as e: logging.warning(f"failed to read {path}: {e}") img = Image.new('RGB', (224, 224), (0, 0, 0)) if self.transform is not None: img = self.transform(img) return img, label这里convert('RGB')很关键。真实数据里混着灰度图、RGBA 图、甚至 CMYK 图，不统一转 RGB，后面 transform 的Normalize会因为通道数不对而报错。这个细节文档里不会强调，但几乎每个真实项目都会遇到。
4.4 DataLoader 配置与多进程调优
数据集写好了，DataLoader的配置同样影响巨大。num_workers设多少？我的经验公式是：CPU 核数的 70% 左右，但要考虑每个 worker 的内存开销。图像任务里每个 worker 可能占几百 MB，设太大直接 OOM。pin_memory=True在 GPU 训练时几乎必开，能加速主机到显存的拷贝。persistent_workers=True在 epoch 较多时能省去反复创建 worker 的开销。
from torch.utils.data import DataLoader train_loader = DataLoader( train_dataset, batch_size=32, shuffle=True, num_workers=8, pin_memory=True, persistent_workers=True, drop_last=True, worker_init_fn=lambda wid: __import__('random').seed(42 + wid) )drop_last=True在分类任务里建议开，避免最后一个不满的 batch 影响 BatchNorm 的统计。worker_init_fn里重新播种是为了让每个 worker 的随机增强序列不同，前面提过这个坑。这些参数单独看都很普通，组合起来才是稳定高效的数据管道。
5. 常见问题与排查技巧实录
5.1 多进程下的典型故障速查
自定义数据集的问题，八成出在多进程上。我整理了一张速查表，覆盖最常见的几种症状和根因。
| 症状 | 可能原因 | 排查与解决 |
|---|---|---|
| 训练卡在第一个 batch 不动 | worker 进程死锁，常见于在__init__里打开了文件句柄或数据库连接 | 把资源打开移到__getitem__，或用worker_init_fn初始化 |
| 内存随 epoch 持续增长 | Dataset 里缓存了数据但没设上限 | 用固定大小的 LRU 缓存，或干脆不缓存 |
| 数据增强每次结果一样 | worker 随机种子相同 | 在worker_init_fn里按 worker id 重新播种 |
| 报 "DataLoader worker exited unexpectedly" | __getitem__抛了未捕获异常 | 在__getitem__里加 try/except 并记录日志 |
| GPU 利用率忽高忽低 | 数据加载是瓶颈 | 增大num_workers、开pin_memory、考虑预加载 |
这张表里的每一条我都在实际项目里遇到过。最坑的是第一条，__init__里打开文件句柄，单进程跑得好好的，一开多进程就死锁，而且报错信息完全不指向根因，查起来非常痛苦。后来养成习惯：__init__里只做纯计算和轻量 IO，任何需要保持打开状态的资源都放到__getitem__里临时打开、用完即关。
5.2 性能瓶颈的定位方法
数据管道慢，先别急着改代码，先定位瓶颈在哪。我的方法是用一个极简的 benchmark：只迭代 DataLoader、不做任何模型计算，看每秒能出多少个 batch。如果这个数字远低于你的预期，问题就在数据侧。
进一步细分：把 transform 关掉再测一次。如果速度大幅提升，瓶颈在 transform，考虑简化增强或者把重操作（如大图缩放）前置到离线阶段。如果关掉 transform 还是慢，瓶颈在 IO，考虑换存储介质、加缓存、或者减少单条数据体积。
还有一个隐蔽的瓶颈：collate_fn。默认 collate 对每个字段做 stack，如果样本里有大张量，stack 本身也要时间。我遇到过一次，单条样本是 4K 图，batch size 16，collate 阶段拷贝数据花了可观的时间。解决办法是在__getitem__里就把图缩到目标尺寸，别把大图传进 collate。
5.3 几个反直觉的实操心得
第一个心得：验证集的 Dataset 不要复用训练集的实例。哪怕 transform 传的是 None，共享实例在多进程下也可能因为状态污染出问题。老老实实建两个实例，代码多两行，省心很多。
第二个心得：索引列表用 tuple 而不是 list。tuple 更省内存，而且在多进程 fork 时拷贝开销更小。数据集几十万条时，这个差异能到几十 MB。
第三个心得：日志要带 worker id。多进程下日志混在一起，不加 worker id 根本分不清是哪个进程出的问题。在__getitem__的异常日志里加上torch.utils.data.get_worker_info()的信息，排查效率翻倍。
第四个心得：先跑通单进程再开多进程。很多人一上来就num_workers=8，出问题后完全不知道是逻辑 bug 还是并发 bug。先用num_workers=0验证逻辑正确，再逐步加 worker，每加一档测一次，问题定位会清晰得多。
6. 进阶扩展：让自定义数据集扛住真实项目
6.1 缓存策略的分层设计
前面提过缓存，这里展开讲一个我常用的分层方案。第一层是内存缓存，用functools.lru_cache或者自己维护一个固定大小的字典，缓存解码后的张量；第二层是磁盘缓存，把预处理后的数据（比如缩放后的图）存成更快的格式，下次直接读；第三层才是原始数据。
分层的关键是缓存键的设计。用文件路径做键最简单，但如果 transform 参数变了，缓存就失效了，必须把 transform 的配置也编进键里。我一般用一个短哈希：hash((path, transform_signature))。这个细节不注意，会出现"改了增强参数但结果没变"的诡异现象，因为读到了旧缓存。
磁盘缓存要注意清理策略。训练中途改了预处理逻辑，旧缓存必须失效，否则数据分布就混了。我的做法是缓存目录名带上配置哈希，配置一变自动用新目录，旧目录手动清理。这个方案土但可靠，比复杂的缓存失效逻辑省心。
6.2 分布式训练下的数据集适配
单机多卡或者多机训练时，数据集要配合DistributedSampler。这里有个容易忽略的点：DistributedSampler会给每个进程分配不同的索引子集，但你的 Dataset 本身不需要感知这件事，它只管按索引取数据。真正需要改的是__len__的语义——在分布式场景下，每个进程看到的"长度"是子集的长度，而不是全集长度。
另一个坑是随机性。分布式训练时，每个进程的随机种子要不同，否则所有卡上的数据增强完全一致，等于变相减小了 batch 的多样性。PyTorch 的DistributedSampler有seed参数，配合set_epoch使用，能保证每个 epoch 的 shuffle 不同且各进程不重复。这套机制要提前设计好，中途改会很麻烦。
6.3 数据质量监控的轻量做法
真实项目里，数据管道跑起来不代表数据是对的。我习惯在 Dataset 里加一个可选的统计钩子：每取 N 条样本，记录一下图像尺寸分布、像素均值、标签分布。这些统计在训练初期就能暴露问题，比如"所有图像都是同一个尺寸说明缩放逻辑写死了"、"某类别占比 90% 说明采样有问题"。
这个钩子要轻量，不能影响性能。我的实现是计数器加定期打印，每 1000 条汇总一次，开销可以忽略。别小看这个习惯，它帮我提前发现过标注文件编码错误导致标签全乱、图像通道顺序搞反导致颜色异常等问题，都是那种不监控就等着模型指标异常才发现的问题。
7. 我个人在实际操作中的体会
写了这么多，最后说几句掏心窝的话。自定义数据集这个事，技术门槛不高，但工程细节极多，属于典型的"会写只要十分钟，写好要踩十个坑"。我现在的习惯是：任何新项目，先花半天把数据管道写扎实、测充分，再动模型。这半天投入的回报率，比调模型超参高得多。
一个具体的建议：给你的 Dataset 写单元测试。不用复杂，就测三件事——len对不对、随机取几条__getitem__返回结构对不对、边界索引（0 和 len-1）不越界。这三条测试能拦住大部分低级错误，而且改代码后跑一遍只要几秒。我见过太多项目因为改了一行索引逻辑，导致训练集和验证集数据串了，指标虚高，最后上线才发现，返工成本巨大。
还有一点：别过度设计。我早期喜欢把 Dataset 写成万能框架，支持各种数据源、各种返回格式，结果代码几百行，自己都记不住逻辑。后来学乖了，一个任务一个 Dataset 类，该重复就重复，清晰比抽象重要。数据管道是基础设施，稳定压倒一切，花哨的抽象在出问题时只会增加排查难度。
