最近在梳理 PyTorch 数据读取机制的时候我把框架训练营第四期里关于 DataLoader 和 Dataset 的内容重新过了一遍顺手把自己实践中的一些笔记也整理了出来。这个东西说复杂不复杂说简单也真不简单——很多人模型结构写得溜一卡到数据加载就开始出各种幺蛾子要么内存爆掉要么一个 epoch 跑半天要么多进程一开直接死锁。其实这些问题大多数不是模型的问题而是对数据读取机制的理解停留在会用层面没弄明白它到底在底层干了什么。这篇文章我会从整体流程讲到源码层面再结合几个真实的踩坑案例把 DataLoader 与 Dataset 的核心机制掰开揉碎讲清楚。无论你是刚入门的新手还是已经跑过几个项目但总被 IO 瓶颈折磨的开发者应该都能从中找到对应的解法。1. 数据读取的整体流程从原始文件到 batch 之间发生了什么很多教程上来就教你怎么定义 Dataset、怎么实例化 DataLoader但我更想先聊聊宏观流程。因为只有先看明白整个链路后面那些参数调起来才不是瞎调。整个数据读取链路大致可以拆成五个环节原始数据存储 - Dataset 逻辑映射 - Sampler 决定顺序 - DataLoader 按 batch 取数 - collate_fn 拼装成张量。这五个环节各自独立又在 DataLoader 的调度下协同工作。第一个环节好理解就是硬盘上的图片、文本、音频文件或内存里的数组。第二个环节 Dataset 负责的是逻辑映射给它一个索引它告诉你这个索引对应的样本长什么样。这里注意它不是把所有数据预先读进内存而是定义了一个按编号取数据的规则。第三个环节很多人会忽略Sampler 决定了每个 epoch 里样本以什么顺序被取出来。shuffleTrue 的时候其实就是在这一层生成一个打乱的索引序列。第四个环节 DataLoader 按 batch_size 去 Sampler 里取一批索引再用这些索引去调用 Dataset 的getitem。第五个环节 collate_fn 把零散的样本堆叠成一个 batch 张量比如把多张图片拼成一个四维张量或者把变长文本补齐到相同长度。这里有个容易混淆的点Dataset 和 DataLoader 的职责边界。Dataset 只负责给定索引返回单个样本它不关心批次、不关心打乱、不关心多线程。DataLoader 才是那个统筹全局的角色它负责迭代、多进程、打乱、批处理、内存锁页等等。你可以把 Dataset 理解成一个数据源说明书DataLoader 是调度中心。我在看训练营课程的时候讲师画了一张数据流向图把整个链路标得非常清楚。后来我实践时发现只要把这一条链路的每一环在脑子里形成画面很多参数设置就不再是死记硬背了。举个例子如果你自定义了一个 Dataset每次getitem都要从磁盘读一张大图再做解码那么即使 DataLoader 设置了 8 个 worker整个链路也会卡在磁盘 IO 和解码上。这时候真正该做的不是堆 worker 数量而是预先做解码缓存或改用内存映射格式。这就是看清完整链路和只会调 API之间的差别。2. Dataset 基类与自定义 Dataset像写说明书一样把数据交给框架PyTorch 的 Dataset 是一个抽象基类它的核心契约就两个方法len和getitem。len返回数据集大小DataLoader 靠它知道一共要迭代多少次getitem接收一个索引返回该索引对应的样本。就这么简单但正是这种简单让它可以适配几乎任何数据形态。自定义 Dataset 的标准写法是继承 torch.utils.data.Dataset然后实现两个方法。比如我有一个存放图片的文件夹图片名从 0.jpg 排到 9999.jpg标签存在 CSV 里我就可以这样写import pandas as pd from torch.utils.data import Dataset from PIL import Image import os class ImageFolderDataset(Dataset): def __init__(self, img_dir, label_csv, transformNone): self.img_dir img_dir self.df pd.read_csv(label_csv) self.transform transform def __len__(self): return len(self.df) def __getitem__(self, index): row self.df.iloc[index] img_path os.path.join(self.img_dir, row[filename]) image Image.open(img_path).convert(RGB) label row[label] if self.transform: image self.transform(image) return image, label这段代码有几个细节值得说。第一init里通常只存路径和元信息不急着把所有图片读进内存。第二getitem每次被调用时实时读取文件这是最通用也最容易的写法代价是每次迭代都有磁盘 IO。第三transform 最好在 getitem 里做而不是提前做好存起来因为这样每次 epoch 都可以做不同的随机增强。不过还有一类 Dataset 叫 IterableDataset适用场景是数据流式到达、无法用索引定位的情况比如从网络流、生成器读取数据或者在超大数据集上做流式训练。它实现的是iter而不是getitem。用 IterableDataset 时要注意多进程下每个 worker 都会独立跑一遍迭代逻辑如果不做切分就会导致重复数据。第三transform 最好在 getitem 里做而不是提前做好存起来因为这样每次 epoch 都可以做不同的随机增强。自定义 Dataset 时最常见的错误是把文件内容一次性全读进内存放在init里。我见过不少新手这么干数据量小的时候感觉很快一旦数据到了几十 GB 就直接 OOM。正确做法是init里只建立索引到文件路径的映射关系真正占用内存的样本在getitem里按需加载。另外getitem的返回值类型要稳定。如果第一个样本返回 (image, label)后面所有样本都要返回这个结构否则 collate_fn 会报错。这也是为什么我建议把数据格式统一写在注释里比如返回类型是 (torch.Tensor[3,H,W], int)方便自己和同事后面看代码。关于 transform 的位置我的习惯是基础预处理比如归一化、尺寸缩放放在 Dataset 内部做掉随机增强比如随机裁剪、翻转也放在 getitem 里做因为这样才能保证每个 epoch 看到的增强都不一样。如果提前把所有增强做完存成新数据集不仅占磁盘也失去了随机增强的意义。3. DataLoader 核心参数逐个拆解batch_size、shuffle、num_workers 背后的真实作用DataLoader 可能是 PyTorch 里参数最多、最容易踩坑的组件之一。很多人的代码就是 DataLoader(dataset, batch_size32, shuffleTrue)跑通了就不再管它一旦要换机器、换数据规模就开始乱调。这一节我把实际发挥作用的参数逐一拆开讲。batch_size 决定了每个 batch 的样本数也决定了显存的占用。它不只是一个数字它直接影响梯度估计的噪声、训练收敛速度和硬件利用率。调的时候有个矛盾点batch 太小一次迭代的梯度方向噪声大训练不稳batch 太大显存可能不够而且容易收敛到尖锐极小值。我的经验是先从显存允许的最大值开始再逐步缩小而不是一开始就设个 32 或 64 了事。shuffle 参数控制每个 epoch 开始时是否打乱样本顺序。为什么需要打乱因为如果数据集天然有序比如按类别排列不打乱的话每个 batch 只包含同一类别的样本梯度会来回震荡收敛极慢。shuffle 的作用是把类别分布打散到各个 batch 里。设置 shuffleTrue 时DataLoader 底层会为每个 epoch 生成一个随机索引序列这个序列来自 Sampler。接着是 num_workers这是最容易让人困惑的参数。它不是 Python 线程而是独立的进程。num_workers0 表示主进程亲自取数据num_workers4 表示额外启动 4 个进程并行调用getitem取数据。这里的关键问题是不是 worker 越多越好。Windows 上 worker 过多还容易爆出 Broken pipe 之类的错误。worker 的数量该怎么定我的经验公式是先看单次取数耗时。如果getitem只做轻量操作比如直接从内存读取 ndarray一个 worker 就够多开反而浪费进程切换的开销。如果getitem要解码大图、读文本、做复杂预处理那么 worker 数量可以从 CPU 核数的一半开始试。训练脚本里通常这样设置from torch.utils.data import DataLoader dataloader DataLoader( dataset, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue, drop_lastFalse, )pin_memory 这个参数第一次见到容易忽略。它的作用是当使用 GPU 训练时将取到的数据放到锁页内存pinned memory里这样从内存拷贝到显存的速度会快很多。代价是锁页内存不能像普通内存那样被换出所以会占用一部分固定内存。显存和数据传输成为瓶颈时pin_memoryTrue 有明显效果如果是 CPU 训练这个参数基本没有意义。drop_last 是控制最后一个 batch 要不要丢弃的参数。假设数据集有 1000 个样本batch_size 设为 96那么最后会剩下 40 个样本凑不成一个完整 batch。drop_lastTrue 就直接丢掉这批样本drop_lastFalse 则会保留然后 collate_fn 可能因为 batch 内样本数不一致而需要特殊处理。还有一个参数 prefetch_factor默认是 2表示每个 worker 预取多少个 batch 的数据。它的作用是隐藏数据加载的延迟当 GPU 在计算第 n 个 batch 时worker 已经在后台准备第 n1、n2 个 batch 的数据了。数据加载太慢但显存还有富余的时候可以适当调大这个值反之如果显存紧张调小可以避免预取的数据占住内存。实际调优时我的建议是先用默认参数跑通然后用观察法定位瓶颈。具体做法是在训练循环里打印每个 epoch 的数据加载耗时和计算耗时。如果加载耗时明显高于计算耗时优先考虑加 num_workers 或 prefetch_factor如果 GPU 利用率持续很低问题往往不在 DataLoader 参数而是数据链路里有重复解码或低效 IO。4. 低频但致命数据加载中的 Dataset 复用与共享缓存设计上一节讲的都是 DataLoader 自带参数但在真实项目里我经常要在 Dataset 内部做一些额外的缓存优化否则单靠调参根本救不回来。先说一个场景我在做一个图像分割项目训练集是 8000 张 2048x2048 的大图。由于显存限制每张图要先切块再送进网络。最初我的 Dataset 在getitem里每次读整张大图、切块、做归一化跑起来一个 epoch 要 15 分钟GPU 利用率只有 30%。定位后发现瓶颈是大量重复的磁盘读取和格式解码。我的解决方案是在 Dataset 里加一个 LRU 缓存把最近访问的原始图缓存下来。因为训练时同一个 epoch 内很多切块来自相同的大图缓存命中率极高。实现并不复杂from collections import OrderedDict class CachedImageDataset(Dataset): def __init__(self, img_list, cache_size16): self.img_list img_list self.cache OrderedDict() self.cache_size cache_size def _read_image(self, img_path): if img_path in self.cache: self.cache.move_to_end(img_path) return self.cache[img_path] img load_image(img_path) if len(self.cache) self.cache_size: self.cache.popitem(lastFalse) self.cache[img_path] img return img def __getitem__(self, index): img_path, crop_box, label self.parse_index(index) img self._read_image(img_path) crop img.crop(crop_box) return transform(crop), label这个方案让 epoch 时间从 15 分钟降到了 6 分钟左右。这里有个关键点缓存的对象是解码后的 PIL Image 或 ndarray而不是原始文件字节。因为解码是 CPU 密集操作缓存解码结果能省下每张图最大的计算开销。但缓存也要小心两个坑。第一个是缓存容量太大时内存占用可能比训练模型本身还高。所以 cache_size 要根据单张图大小和内存总量来估算。第二个坑是缓存只对本地读取有效如果数据源在云端或网络文件系统加缓存反而可能造成一致性问题。另外一个常见但很少被提及的设计是多个 Datasets 共享同一份缓存。例如训练集和验证集都包含部分相同的基础图片只是标注不同。如果分别创建独立的 Dataset每份缓存都是重复的如果让两个 Dataset 引用同一个缓存对象内存占用直接减半。这类缓存设计的核心思想是Dataset 的价值不只是把文件路径变成样本你可以在它内部做任何合理的预处理和缓存逻辑只要保证getitem稳定返回即可。训练营里讲的是标准写法实际落地时一定要结合数据形态做定制。5. 进阶Sampler、collate_fn 与分布式场景下的数据读取设计如果把前三节的内容比作 DataLoader 的常规武器那么 Sampler、collate_fn 和分布式采样就是非常规武器。平时不一定用得上但遇到特殊需求时它们能解决大问题。先讲 Sampler。DataLoader 里有一个 sampler 参数它决定了样本的取用顺序。前面说过 shuffleTrue 时会自动创建一个 RandomSampler其实我们完全可以自己传入 Sampler 来实现更复杂的逻辑。常见的两个场景是类别不平衡采样WeightedRandomSampler和顺序采样SequentialSampler。WeightedRandomSampler 非常适合处理极度不平衡的分类问题。比如有一个二分类任务正样本只有 1%负样本 99%如果使用普通随机采样很多 batch 里几乎全是负样本。WeightedRandomSampler 会给每个样本分配一个权重权重越高的样本被抽中的概率越大从而平衡批内的类别比例。用法如下from torch.utils.data import WeightedRandomSampler weights [0.5 if label 1 else 0.01 for label in all_labels] sampler WeightedRandomSampler(weights, num_sampleslen(weights), replacementTrue) dataloader DataLoader(dataset, batch_size32, samplersampler)注意用了 sampler 之后shuffle 参数不能再设 True因为采样顺序完全由 sampler 决定shuffle 参数本身会创建另一个 sampler两者冲突。再说 collate_fn。默认情况下DataLoader 调用默认的 collate 逻辑它假设所有样本都是相同尺寸的张量然后把它们堆叠成一个更高维的张量。如果样本是变长序列、不同尺寸的图、或者包含字符串标签默认逻辑就会报错。这时你需要自定义 collate_fn。经典案例是自然语言处理里的变长序列。假设一句话长度是 10 个 token另一句话是 25 个 token默认 collate 无法堆叠所以我们要在 collate_fn 里做 paddingdef collate_batch(batch): texts, labels zip(*batch) max_len max(len(t) for t in texts) padded_texts [t [0] * (max_len - len(t)) for t in texts] return torch.tensor(padded_texts), torch.tensor(labels) dataloader DataLoader(dataset, batch_size16, collate_fncollate_batch)collate_fn 的另一个用途是过滤异常样本。比如某个样本在读取时损坏了你可以在 collate_fn 里丢掉它返回一个正常的 batch。注意这样会让 batch 大小变小需要配合 drop_last 或动态调整。分布式训练是数据读取的另一个大坑。当你用 DistributedDataParallel 在 8 张卡上训练时如果直接用普通的 DataLoader每张卡都会拿到全部数据的采样结果等价于每个 batch 数据重复了 8 份。正确的做法是用 DistributedSampler它会根据总进程数把数据切成多份保证每个进程看到的是互不重叠的数据子集from torch.utils.data.distributed import DistributedSampler sampler DistributedSampler(dataset, num_replicasworld_size, rankrank) dataloader DataLoader(dataset, batch_size32, samplersampler)还有一个细节DistributedSampler 在每个 epoch 开始前需要调用 set_epoch(epoch) 来重新打乱数据顺序否则每个 epoch 的划分顺序都是一样的模型会看到固定的数据分组影响训练效果。我在实践中发现分布式训练下 num_workers 的设置也要保守一点因为每张卡上都开了多个 worker总的进程数会非常惊人容易把机器内存吃满。我一般会按每卡 2-4 个 worker 起步观察系统负载再往上加。6. 最后分享几点我在实操里的经验心得关于数据读取这套机制我想最后再分享几个用真金白银换来的经验。第一遇到训练慢先别急着上多进程先确定瓶颈是磁盘 IO、解码还是 GPU 传输。最简单的方法做一次纯加载测试跑一个 epoch 不训练只打印耗时。如果这个测试时间就很长问题在数据读取链路如果测试很快但实际训练慢问题可能在 GPU 传输或预处理。第二数据读取的优化要顺着数据形态来。小数据全部放进内存是最快的方案中等规模数据可以考虑内存映射或者缓存解码结果超大规模数据就要靠流式读取和预处理管道了。没有银弹只有匹配。第三自定义 Dataset 时一定要保持干净的分层路径管理放最外层样本读取放中间层数据增强放内层。这样做的好处是出了问题能快速定位到具体环节而不是在几百行代码里到处找。第四务必关注 worker 进程的异常处理。比如在 Windows 上如果 Dataset 定义写在主脚本里而不是独立模块里多进程时可能会出现无限递归或重复执行的问题。业界通常的解法是把 Dataset 和 DataLoader 相关的类放在单独的 .py 文件里再用 from module import 的方式引入这样多进程重新导入时不会发生诡异行为。我还想强调一下打印调试的重要性。在 Dataset 的getitem里加一个 print很多人觉得这会拖慢速度但如果只在训练的前 5 步打印完全不影响性能却能让你立刻确认数据读取是否正确。类似的自定义 collate_fn 出来后可以先构造一个小样本集测试一下形状是否正确再上全量数据能省很多排查时间。数据读取机制在整个 PyTorch 生态里属于入门门槛低、但优化空间极大的部分。像我一开始那样只会调 DataLoader 参数到后来开始自己写缓存、自定义 Sampler、处理分布式采样每一步都能感觉到训练效率的变化。希望这篇笔记能帮你少走一些弯路。