1. 为什么 dataloader 值得单独拎出来讲搞过深度学习项目的人都有一个共识模型结构再漂亮如果数据管道没搭好训练过程就是一场灾难。我见过太多人把 ResNet、Transformer 玩得飞起结果卡在数据加载上——GPU 利用率常年 30% 以下训练一个 epoch 要等半天排查半天发现是 dataloader 的 num_workers 设成了 0。torch.utils.data.DataLoader这个类表面上看 API 简单到不行无非就是传个 dataset、设个 batch_size、开几个 worker。但真正用起来里面的坑一个接一个多进程 worker 的内存爆炸、自定义 collate_fn 写错导致维度对不上、shuffle 和 sampler 冲突直接报错、pin_memory 开了反而更慢……这些问题在官方文档里往往一笔带过只有实际踩过才知道疼。这篇文章面向的是已经能跑通简单训练脚本、但想真正把数据管道吃透的人。不管你是做图像分类、目标检测、NLP 还是推荐系统只要用 PyTorchdataloader 就是绕不开的基础设施。我会从设计思路讲到实操细节从参数含义讲到排查技巧尽量把我在实际项目里积累的经验都倒出来。读完你至少能做到两件事第一知道每个参数背后的取舍逻辑第二遇到数据加载相关的报错和性能问题能自己定位并解决。2. DataLoader 的整体设计与核心思路拆解2.1 它到底解决了什么问题在没有 DataLoader 的世界里训练循环大概长这样你有一个巨大的数据集存在硬盘上每次取一个 batch 就得手动写循环去读文件、做预处理、拼成 tensor、搬到 GPU。代码又臭又长不说还特别容易出错。更关键的是数据读取和模型计算是串行的——GPU 在算的时候 CPU 闲着CPU 在读数据的时候 GPU 闲着两边互相等效率极低。DataLoader 的核心价值就两个词抽象和并行。抽象是指它把“从数据集里取一批样本”这件事封装成了一个可迭代对象你只需要关心 dataset 怎么定义剩下的批次组装、打乱顺序、多进程调度它全帮你管了。并行是指它通过多进程机制让数据读取和预处理在 CPU 侧并行进行同时主进程可以持续向 GPU 喂数据实现计算和加载的重叠。打个比方Dataset 是你家冰箱里面放着各种食材样本DataLoader 是一个配菜师傅团队你告诉他“每道菜要哪些食材、怎么切”他就安排人从冰箱里取料、切好、摆盘然后一批一批端给你。你只管炒菜训练模型不用自己去冰箱里翻。2.2 Dataset 和 DataLoader 的分工边界很多人初学的时候会把这两者搞混。记住一个原则Dataset 负责“单个样本怎么取”DataLoader 负责“一批样本怎么组”。Dataset 的核心是__getitem__和__len__两个方法。__len__告诉 DataLoader 一共有多少样本__getitem__告诉它第 i 个样本长什么样。至于这个样本是从内存里拿、从硬盘读、还是从数据库查Dataset 不关心DataLoader 也不关心。DataLoader 拿到 Dataset 之后做的事情是决定这一轮按什么顺序取样本shuffle 或 sampler、决定一次取多少个batch_size、决定用几个进程去取num_workers、决定取出来之后怎么拼成一个 batchcollate_fn、决定要不要提前搬到固定内存pin_memory。这个分工的好处是解耦。你可以写一个通用的 Dataset 处理某种数据格式然后用不同的 DataLoader 配置来适配训练、验证、测试等不同场景。比如训练时 shuffleTrue、num_workers8验证时 shuffleFalse、num_workers4Dataset 代码完全不用改。2.3 多进程加载的架构逻辑DataLoader 的多进程机制是它最核心也最容易出问题的地方。当你设置 num_workers 0 时PyTorch 会启动指定数量的子进程每个子进程持有一份 Dataset 的副本在 Linux 上通过 fork 创建在 Windows 上通过 spawn。主进程负责调度和组装 batch子进程负责实际的数据读取和预处理。具体流程是这样的主进程把“需要取哪些索引”的任务分发给各个 worker每个 worker 根据自己的 Dataset 副本去取样本、做 transform然后把处理好的样本放到共享队列里。主进程从队列里取数据用 collate_fn 拼成 batch返回给训练循环。这里有个关键点每个 worker 都会复制一份 Dataset 对象。如果你的 Dataset 在__init__里加载了很大的数据比如把所有图片读进内存那么 num_workers8 就意味着内存里会有 8 份副本直接爆炸。这是新手最容易踩的坑之一。正确的做法是把大数据放在 Dataset 外部通过共享内存或内存映射的方式访问或者用__getitem__懒加载。另一个关键点是随机种子的处理。PyTorch 会为每个 worker 设置不同的随机种子确保数据增强的随机性不会因为多进程而重复。但这个机制依赖于你正确使用了 PyTorch 的随机数生成器。如果你在 Dataset 里用了 Python 原生的random模块可能会发现所有 worker 产生的随机变换一模一样。解决办法是在 Dataset 的__init__里用torch.initial_seed()或者手动为每个 worker 设置种子。3. 核心参数逐个拆解与实操要点3.1 batch_size 与 drop_last 的取舍batch_size是最直观的参数但它的选择不只是“显存够不够”这么简单。batch_size 影响的不只是显存占用还有梯度估计的方差、BatchNorm 的统计量质量、训练收敛的速度。从显存角度算一个 batch 的显存占用大致是batch_size × 单样本显存 × (1 梯度 优化器状态系数)。对于 Adam 优化器优化器状态通常是参数量的 2 倍所以总系数大概是 4 左右。你可以先用 batch_size1 跑一下看显存占用然后线性估算能开多大。但显存不是唯一约束。batch_size 太小比如 1 或 2BatchNorm 的统计量会非常不稳定梯度噪声也大训练可能不收敛。batch_size 太大梯度估计太准反而可能陷入 sharp minimum泛化性能下降。实践中图像分类任务常用 32 到 256检测任务因为单样本显存大常用 2 到 16。drop_last这个参数容易被忽略。默认是 False意味着最后一个不完整的 batch 也会被返回。这在训练时通常没问题但在某些场景下会出问题比如你用 BatchNorm最后一个 batch 只有一个样本统计量会非常离谱或者你的 loss 计算依赖 batch 内样本间的交互比如对比学习不完整的 batch 会导致逻辑错误。所以训练时我一般建议设drop_lastTrue验证时设 False因为验证不需要反向传播不完整 batch 无所谓。3.2 shuffle、sampler 与 batch_sampler 的关系shuffleTrue是最常用的打乱方式但它的底层其实是创建了一个RandomSampler。这里有个规则shuffle、sampler、batch_sampler 三者互斥。你设了 sampler就不能设 shuffle设了 batch_sampler就不能设 batch_size、shuffle、sampler、drop_last。为什么要有 sampler 这个抽象因为有些场景下打乱方式不是简单的随机。比如分布式训练时需要用DistributedSampler把数据切分到各个 GPU 上每个 GPU 只看到一部分数据。类别不平衡时需要用WeightedRandomSampler按类别权重采样让少数类被更频繁地抽到。序列任务中可能需要按长度分组采样BucketIterator的思路减少 padding 浪费。WeightedRandomSampler的用法值得展开说一下。假设你有 1000 个样本类别 A 有 900 个类别 B 有 100 个。直接随机采样的话每个 batch 里 B 类样本很少模型会偏向 A。用 WeightedRandomSampler给每个样本一个权重A 类样本权重设为 1/900B 类设为 1/100这样采样时 B 类被抽到的概率就大大提高了。注意num_samples参数要设成你想要的 epoch 长度通常等于数据集大小。3.3 num_workers 设多少才合适这是被问得最多的问题没有之一。答案不是固定的取决于你的数据加载瓶颈在哪。如果数据预处理很轻比如数据已经在内存里只需要简单切片num_workers 设 0 或 1 就够了开多了反而因为进程间通信开销而变慢。如果数据预处理很重比如从硬盘读大图、做复杂的增强num_workers 可以设到 CPU 核心数甚至更多。我的一般建议是从 4 开始试逐步增加到 8、16观察 GPU 利用率和每个 epoch 的耗时。如果增加 worker 数量后 GPU 利用率不再提升说明瓶颈已经不在数据加载了再加也没用。如果增加后反而变慢说明进程间通信或内存拷贝成了新瓶颈。还有一个经验值num_workers 不要超过 CPU 物理核心数。比如你的机器是 8 核 16 线程num_workers 设到 8 左右比较合适设到 16 会因为超线程竞争而效率下降。另外如果 num_workers 0建议在训练脚本的入口加上if __name__ __main__:保护否则在 Windows 上会无限递归创建进程。3.4 pin_memory 与 prefetch_factor 的配合pin_memoryTrue的作用是把 CPU 内存中的数据提前拷贝到“固定内存”pinned memory中这样从 CPU 到 GPU 的传输可以用 DMA 直接进行速度更快。这个参数在 GPU 训练时几乎总是建议开启代价是占用一些额外的 CPU 内存。但 pin_memory 不是万能的。如果你的数据量很小或者 GPU 计算本身就是瓶颈pin_memory 带来的收益微乎其微。另外pin_memory 只对 CUDA 张量有效如果你用的是 CPU 训练或者 MPS 后端开了也没用。prefetch_factor控制每个 worker 提前取多少个 batch 的数据。默认是 2意味着每个 worker 会提前准备好 2 个 batch 放在队列里。如果你发现 GPU 经常“等数据”可以适当增大这个值比如设到 4 或 6。但设太大也会占用更多内存而且如果数据加载本身很快设大了反而增加延迟。这两个参数的关系是pin_memory 优化的是“传输速度”prefetch_factor 优化的是“供给连续性”。两者配合好了才能让 GPU 始终有活干。4. 自定义 Dataset 与 collate_fn 的实操细节4.1 写一个健壮的 Dataset 类一个标准的 Dataset 类长这样from torch.utils.data import Dataset class MyDataset(Dataset): def __init__(self, file_list, transformNone): self.file_list file_list self.transform transform def __len__(self): return len(self.file_list) def __getitem__(self, idx): path self.file_list[idx] # 读取数据 data self._load_data(path) # 应用变换 if self.transform: data self.transform(data) return data, label看起来简单但有几个细节决定了它是否健壮。第一__getitem__里不要做太重的初始化操作。比如每次调用都重新打开一个数据库连接、重新加载一个预训练模型这会严重拖慢速度。应该在__init__里做好这些准备工作__getitem__只做轻量的读取和变换。第二异常处理要到位。如果某个样本文件损坏了__getitem__抛异常会导致整个训练崩溃。可以在__getitem__里加 try-except遇到坏样本就返回一个替代样本或者随机跳到另一个索引。但要注意如果坏样本太多这种“跳过”策略会掩盖数据质量问题最好还是先清洗数据。第三注意索引的类型。__getitem__接收的 idx 可能是 Python int也可能是 numpy int64 或 torch tensor。如果你用 idx 去索引一个 Python listnumpy int64 是可以的但 torch tensor 不行。保险的做法是在开头加一句idx int(idx)。4.2 collate_fn 的默认行为与自定义场景默认的collate_fn做的事情很简单把 batch 个样本的 tuple 按位置拆开每个位置上的元素拼成一个 batch。比如每个样本返回(image, label)默认 collate 会返回(batch_images, batch_labels)其中 batch_images 是把所有 image 沿第 0 维堆叠。但很多场景下默认行为不够用。比如目标检测中每张图片的标注框数量不一样没法直接堆叠需要自定义 collate_fn 把标注整理成列表。NLP 中每个句子的长度不一样需要 padding 到同一长度同时生成 attention mask。多模态任务中不同模态的数据需要分别处理可能还要做对齐。写自定义 collate_fn 的时候最容易犯的错误是维度搞混。建议在函数开头打印一下输入的形状确认无误后再写后续逻辑。另外collate_fn 是在主进程里执行的不是在 worker 里所以如果里面有耗时操作会阻塞整个数据流。尽量把重活放在 Dataset 的__getitem__里collate_fn 只做组装。4.3 数据增强与 transform 的放置位置数据增强应该放在 Dataset 的__getitem__里还是放在 collate_fn 里答案是绝大多数情况下放在__getitem__里。原因是__getitem__是在 worker 进程中执行的多个 worker 可以并行做增强充分利用 CPU 多核。如果放在 collate_fn 里就是在主进程串行执行速度会慢很多。但有一个例外如果增强操作需要 batch 内样本间的交互比如 MixUp、CutMix那就必须放在 collate_fn 里因为__getitem__一次只能看到一个样本。另外验证集和测试集的 transform 通常和训练集不一样。训练集用随机裁剪、随机翻转等增强验证集用中心裁剪、不翻转。这个切换可以通过在 Dataset 初始化时传入不同的 transform 来实现也可以通过dataset.transform val_transform动态修改。我一般倾向于前者逻辑更清晰。5. 常见问题与排查技巧实录5.1 多进程相关的典型报错报错一RuntimeError: DataLoader worker (pid xxx) is killed by signal: Bus error这个报错通常是因为共享内存不够。Linux 系统默认的/dev/shm大小一般是内存的一半如果你开了很多 worker每个 worker 又往共享内存里写大量数据就会撑爆。解决办法有两个一是减小 num_workers二是增大/dev/shm的大小需要系统权限。报错二TypeError: cannot pickle _thread.lock object这个报错是因为 Dataset 里包含了不可序列化的对象比如数据库连接、文件句柄、线程锁。在 Windows 上worker 是通过 spawn 创建的需要把 Dataset 序列化后传给子进程所以要求 Dataset 必须可 pickle。解决办法是把这些不可序列化的对象放在__getitem__里懒加载或者用__getstate__和__setstate__自定义序列化行为。报错三RuntimeError: Cowardly refusing to serialize non-leaf tensor which requires_grad这个报错是因为 Dataset 返回了需要梯度的 tensor。DataLoader 在 worker 和主进程之间传递数据时会尝试序列化 tensor而需要梯度的 tensor 不能被序列化。解决办法是在 Dataset 里用torch.no_grad()包裹或者用.detach()去掉梯度。5.2 性能问题的排查思路数据加载慢怎么定位瓶颈我的排查顺序是这样的第一步把 num_workers 设为 0测一下纯数据加载的耗时。如果 num_workers0 时每个 epoch 要 100 秒num_workers8 时要 20 秒说明多进程确实起了作用但还可以继续优化。第二步用time.time()在__getitem__里打点看看时间花在哪里。是文件 IO 慢还是 transform 慢还是其他操作慢。如果是文件 IO 慢考虑把数据转成更高效的格式比如 LMDB、HDF5、WebDataset如果是 transform 慢考虑用 GPU 做增强比如 NVIDIA DALI或者用更快的库比如 OpenCV 替代 PIL。第三步检查 GPU 利用率。用nvidia-smi -l 1或者watch -n 1 nvidia-smi观察。如果 GPU 利用率经常掉到 0说明数据供给跟不上如果 GPU 利用率稳定在 90% 以上说明数据加载不是瓶颈不用再折腾了。第四步检查内存和共享内存使用。用free -h看内存用df -h /dev/shm看共享内存。如果共享内存快满了减小 num_workers 或者减小 prefetch_factor。5.3 常见问题速查表问题现象可能原因解决办法GPU 利用率低训练慢num_workers 太小或为 0增大 num_workers开启 pin_memory内存爆炸Dataset 在__init__里加载了大数据每个 worker 一份副本改为懒加载或用内存映射报 Bus error共享内存不足减小 num_workers 或增大 /dev/shm报 pickle 错误Dataset 含不可序列化对象懒加载或自定义序列化随机增强失效用了 Python 原生 random改用 torch 的随机函数最后一个 batch 报错drop_lastFalse 且 BatchNorm 遇到单样本设 drop_lastTrueshuffle 和 sampler 冲突同时设了多个互斥参数只保留一个Windows 上无限递归没加if __name__ __main__加上入口保护5.4 几个我踩过的坑第一个坑在 Dataset 的__init__里做数据过滤。我曾经写过一个 Dataset在__init__里遍历所有文件把损坏的文件剔除掉。结果 num_workers8 时每个 worker 都做了一遍过滤启动时间翻了 8 倍。后来改成在外部做好过滤把干净的文件列表传进来启动时间立刻降下来了。第二个坑在__getitem__里用random.random()做随机增强。结果发现所有 worker 产生的随机数序列一模一样因为每个 worker 的 Python random 种子相同。改成torch.rand()之后问题解决因为 PyTorch 会为每个 worker 设置不同的种子。第三个坑pin_memoryTrue 但 num_workers0。这种情况下 pin_memory 是在主进程里做的会阻塞训练循环反而更慢。pin_memory 只有在 num_workers 0 时才有意义因为这时数据拷贝可以和主进程的计算重叠。第四个坑用IterableDataset时忘了设batch_sizeNone。IterableDataset是自己控制 batch 的如果 DataLoader 再设 batch_size会报错。正确做法是设batch_sizeNone让 Dataset 自己 yield batch。6. 进阶用法与性能调优6.1 IterableDataset 的适用场景前面讲的都是MapStyleDataset也就是通过索引访问的 Dataset。PyTorch 还提供了IterableDataset它不通过索引访问而是直接 yield 样本。适合以下场景数据流是无限的或者非常大的没法预先知道长度。数据需要按顺序读取比如从数据库游标、日志流、网络流中读取。数据分片存储在多个文件中需要动态决定读取顺序。IterableDataset的多进程处理和MapStyleDataset不一样。因为每个 worker 都会独立地迭代数据如果不做特殊处理多个 worker 会读到重复的数据。解决办法是用worker_init_fn给每个 worker 分配不同的数据分片或者用torch.utils.data.get_worker_info()在运行时获取 worker 编号据此决定读哪些数据。6.2 用 Sampler 实现分布式训练的数据切分分布式训练时每个 GPU 应该只看到数据集的一部分同时保证每个 epoch 所有 GPU 合起来能看到全部数据。DistributedSampler就是干这个的。它的工作原理是根据当前进程的 rank 和总进程数把数据集分成若干份每个进程只取自己那一份。同时每个 epoch 会重新打乱顺序保证不同 epoch 看到的数据组合不同。使用时要注意DistributedSampler需要设置shuffleTrue来打乱同时 DataLoader 的 shuffle 要设为 False。另外每个 epoch 开始前要调用sampler.set_epoch(epoch)否则每个 epoch 的打乱顺序会一样。6.3 数据预取与缓存策略如果数据预处理很耗时可以考虑把处理好的数据缓存起来。最简单的做法是在__getitem__里加一个内存缓存字典第一次读取时做预处理并缓存后续直接返回缓存结果。但要注意缓存大小别把内存撑爆。更高级的做法是用磁盘缓存。比如把预处理后的数据存成 numpy 数组或 tensor 文件下次直接读。或者用 LMDB 这样的嵌入式数据库支持高效的随机读取。还有一种策略是预取。除了 DataLoader 自带的 prefetch_factor你还可以在 Dataset 层面做预取。比如在__getitem__里提前读取下一个样本的数据利用 IO 等待的时间做其他事情。但这会增加代码复杂度一般只在极端性能要求下才用。6.4 性能调优的检查清单调优数据管道时我一般按这个清单逐项检查num_workers 是否设到了合适值观察 GPU 利用率pin_memory 是否开启GPU 训练时prefetch_factor 是否足够观察 GPU 是否有等待Dataset 是否有不必要的重复初始化transform 是否用了高效的库OpenCV 通常比 PIL 快数据格式是否高效避免大量小文件考虑打包格式是否有内存泄漏观察长时间训练的内存变化collate_fn 是否有耗时操作应该尽量轻量7. 一些个人体会DataLoader 这个组件刚入门的时候觉得很简单用久了才发现水深。它的每一个参数背后都有设计者的权衡每一个报错背后都有特定的运行机制。我自己的经验是不要等到出问题了才去查文档而是提前把多进程、内存、序列化这些机制搞清楚后面能省很多时间。另外数据管道的优化是有上限的。当 GPU 利用率稳定在 90% 以上时再优化数据加载的收益就很小了这时候应该把精力放在模型结构、学习率调度、正则化等更影响最终效果的地方。不要陷入“为了优化而优化”的陷阱。最后分享一个习惯每次搭好一个新的数据管道我会先用小数据量跑几个 epoch确认数据读取正确、增强效果符合预期、batch 组装没问题然后再上全量数据。这个习惯帮我避免了很多“训练了半天才发现数据有问题”的尴尬。 
长河君 · 2026/10/22 16:11:30 · 阅读 1.2k PyTorch DataLoader 深度解析：多进程加载、参数调优与性能排查实战
1. 为什么 dataloader 值得单独拎出来讲搞过深度学习项目的人都有一个共识模型结构再漂亮如果数据管道没搭好训练过程就是一场灾难。我见过太多人把 ResNet、Transformer 玩得飞起结果卡在数据加载上——GPU 利用率常年 30% 以下训练一个 epoch 要等半天排查半天发现是 dataloader 的 num_workers 设成了 0。torch.utils.data.DataLoader这个类表面上看 API 简单到不行无非就是传个 dataset、设个 batch_size、开几个 worker。但真正用起来里面的坑一个接一个多进程 worker 的内存爆炸、自定义 collate_fn 写错导致维度对不上、shuffle 和 sampler 冲突直接报错、pin_memory 开了反而更慢……这些问题在官方文档里往往一笔带过只有实际踩过才知道疼。这篇文章面向的是已经能跑通简单训练脚本、但想真正把数据管道吃透的人。不管你是做图像分类、目标检测、NLP 还是推荐系统只要用 PyTorchdataloader 就是绕不开的基础设施。我会从设计思路讲到实操细节从参数含义讲到排查技巧尽量把我在实际项目里积累的经验都倒出来。读完你至少能做到两件事第一知道每个参数背后的取舍逻辑第二遇到数据加载相关的报错和性能问题能自己定位并解决。2. DataLoader 的整体设计与核心思路拆解2.1 它到底解决了什么问题在没有 DataLoader 的世界里训练循环大概长这样你有一个巨大的数据集存在硬盘上每次取一个 batch 就得手动写循环去读文件、做预处理、拼成 tensor、搬到 GPU。代码又臭又长不说还特别容易出错。更关键的是数据读取和模型计算是串行的——GPU 在算的时候 CPU 闲着CPU 在读数据的时候 GPU 闲着两边互相等效率极低。DataLoader 的核心价值就两个词抽象和并行。抽象是指它把“从数据集里取一批样本”这件事封装成了一个可迭代对象你只需要关心 dataset 怎么定义剩下的批次组装、打乱顺序、多进程调度它全帮你管了。并行是指它通过多进程机制让数据读取和预处理在 CPU 侧并行进行同时主进程可以持续向 GPU 喂数据实现计算和加载的重叠。打个比方Dataset 是你家冰箱里面放着各种食材样本DataLoader 是一个配菜师傅团队你告诉他“每道菜要哪些食材、怎么切”他就安排人从冰箱里取料、切好、摆盘然后一批一批端给你。你只管炒菜训练模型不用自己去冰箱里翻。2.2 Dataset 和 DataLoader 的分工边界很多人初学的时候会把这两者搞混。记住一个原则Dataset 负责“单个样本怎么取”DataLoader 负责“一批样本怎么组”。Dataset 的核心是__getitem__和__len__两个方法。__len__告诉 DataLoader 一共有多少样本__getitem__告诉它第 i 个样本长什么样。至于这个样本是从内存里拿、从硬盘读、还是从数据库查Dataset 不关心DataLoader 也不关心。DataLoader 拿到 Dataset 之后做的事情是决定这一轮按什么顺序取样本shuffle 或 sampler、决定一次取多少个batch_size、决定用几个进程去取num_workers、决定取出来之后怎么拼成一个 batchcollate_fn、决定要不要提前搬到固定内存pin_memory。这个分工的好处是解耦。你可以写一个通用的 Dataset 处理某种数据格式然后用不同的 DataLoader 配置来适配训练、验证、测试等不同场景。比如训练时 shuffleTrue、num_workers8验证时 shuffleFalse、num_workers4Dataset 代码完全不用改。2.3 多进程加载的架构逻辑DataLoader 的多进程机制是它最核心也最容易出问题的地方。当你设置 num_workers 0 时PyTorch 会启动指定数量的子进程每个子进程持有一份 Dataset 的副本在 Linux 上通过 fork 创建在 Windows 上通过 spawn。主进程负责调度和组装 batch子进程负责实际的数据读取和预处理。具体流程是这样的主进程把“需要取哪些索引”的任务分发给各个 worker每个 worker 根据自己的 Dataset 副本去取样本、做 transform然后把处理好的样本放到共享队列里。主进程从队列里取数据用 collate_fn 拼成 batch返回给训练循环。这里有个关键点每个 worker 都会复制一份 Dataset 对象。如果你的 Dataset 在__init__里加载了很大的数据比如把所有图片读进内存那么 num_workers8 就意味着内存里会有 8 份副本直接爆炸。这是新手最容易踩的坑之一。正确的做法是把大数据放在 Dataset 外部通过共享内存或内存映射的方式访问或者用__getitem__懒加载。另一个关键点是随机种子的处理。PyTorch 会为每个 worker 设置不同的随机种子确保数据增强的随机性不会因为多进程而重复。但这个机制依赖于你正确使用了 PyTorch 的随机数生成器。如果你在 Dataset 里用了 Python 原生的random模块可能会发现所有 worker 产生的随机变换一模一样。解决办法是在 Dataset 的__init__里用torch.initial_seed()或者手动为每个 worker 设置种子。3. 核心参数逐个拆解与实操要点3.1 batch_size 与 drop_last 的取舍batch_size是最直观的参数但它的选择不只是“显存够不够”这么简单。batch_size 影响的不只是显存占用还有梯度估计的方差、BatchNorm 的统计量质量、训练收敛的速度。从显存角度算一个 batch 的显存占用大致是batch_size × 单样本显存 × (1 梯度 优化器状态系数)。对于 Adam 优化器优化器状态通常是参数量的 2 倍所以总系数大概是 4 左右。你可以先用 batch_size1 跑一下看显存占用然后线性估算能开多大。但显存不是唯一约束。batch_size 太小比如 1 或 2BatchNorm 的统计量会非常不稳定梯度噪声也大训练可能不收敛。batch_size 太大梯度估计太准反而可能陷入 sharp minimum泛化性能下降。实践中图像分类任务常用 32 到 256检测任务因为单样本显存大常用 2 到 16。drop_last这个参数容易被忽略。默认是 False意味着最后一个不完整的 batch 也会被返回。这在训练时通常没问题但在某些场景下会出问题比如你用 BatchNorm最后一个 batch 只有一个样本统计量会非常离谱或者你的 loss 计算依赖 batch 内样本间的交互比如对比学习不完整的 batch 会导致逻辑错误。所以训练时我一般建议设drop_lastTrue验证时设 False因为验证不需要反向传播不完整 batch 无所谓。3.2 shuffle、sampler 与 batch_sampler 的关系shuffleTrue是最常用的打乱方式但它的底层其实是创建了一个RandomSampler。这里有个规则shuffle、sampler、batch_sampler 三者互斥。你设了 sampler就不能设 shuffle设了 batch_sampler就不能设 batch_size、shuffle、sampler、drop_last。为什么要有 sampler 这个抽象因为有些场景下打乱方式不是简单的随机。比如分布式训练时需要用DistributedSampler把数据切分到各个 GPU 上每个 GPU 只看到一部分数据。类别不平衡时需要用WeightedRandomSampler按类别权重采样让少数类被更频繁地抽到。序列任务中可能需要按长度分组采样BucketIterator的思路减少 padding 浪费。WeightedRandomSampler的用法值得展开说一下。假设你有 1000 个样本类别 A 有 900 个类别 B 有 100 个。直接随机采样的话每个 batch 里 B 类样本很少模型会偏向 A。用 WeightedRandomSampler给每个样本一个权重A 类样本权重设为 1/900B 类设为 1/100这样采样时 B 类被抽到的概率就大大提高了。注意num_samples参数要设成你想要的 epoch 长度通常等于数据集大小。3.3 num_workers 设多少才合适这是被问得最多的问题没有之一。答案不是固定的取决于你的数据加载瓶颈在哪。如果数据预处理很轻比如数据已经在内存里只需要简单切片num_workers 设 0 或 1 就够了开多了反而因为进程间通信开销而变慢。如果数据预处理很重比如从硬盘读大图、做复杂的增强num_workers 可以设到 CPU 核心数甚至更多。我的一般建议是从 4 开始试逐步增加到 8、16观察 GPU 利用率和每个 epoch 的耗时。如果增加 worker 数量后 GPU 利用率不再提升说明瓶颈已经不在数据加载了再加也没用。如果增加后反而变慢说明进程间通信或内存拷贝成了新瓶颈。还有一个经验值num_workers 不要超过 CPU 物理核心数。比如你的机器是 8 核 16 线程num_workers 设到 8 左右比较合适设到 16 会因为超线程竞争而效率下降。另外如果 num_workers 0建议在训练脚本的入口加上if __name__ __main__:保护否则在 Windows 上会无限递归创建进程。3.4 pin_memory 与 prefetch_factor 的配合pin_memoryTrue的作用是把 CPU 内存中的数据提前拷贝到“固定内存”pinned memory中这样从 CPU 到 GPU 的传输可以用 DMA 直接进行速度更快。这个参数在 GPU 训练时几乎总是建议开启代价是占用一些额外的 CPU 内存。但 pin_memory 不是万能的。如果你的数据量很小或者 GPU 计算本身就是瓶颈pin_memory 带来的收益微乎其微。另外pin_memory 只对 CUDA 张量有效如果你用的是 CPU 训练或者 MPS 后端开了也没用。prefetch_factor控制每个 worker 提前取多少个 batch 的数据。默认是 2意味着每个 worker 会提前准备好 2 个 batch 放在队列里。如果你发现 GPU 经常“等数据”可以适当增大这个值比如设到 4 或 6。但设太大也会占用更多内存而且如果数据加载本身很快设大了反而增加延迟。这两个参数的关系是pin_memory 优化的是“传输速度”prefetch_factor 优化的是“供给连续性”。两者配合好了才能让 GPU 始终有活干。4. 自定义 Dataset 与 collate_fn 的实操细节4.1 写一个健壮的 Dataset 类一个标准的 Dataset 类长这样from torch.utils.data import Dataset class MyDataset(Dataset): def __init__(self, file_list, transformNone): self.file_list file_list self.transform transform def __len__(self): return len(self.file_list) def __getitem__(self, idx): path self.file_list[idx] # 读取数据 data self._load_data(path) # 应用变换 if self.transform: data self.transform(data) return data, label看起来简单但有几个细节决定了它是否健壮。第一__getitem__里不要做太重的初始化操作。比如每次调用都重新打开一个数据库连接、重新加载一个预训练模型这会严重拖慢速度。应该在__init__里做好这些准备工作__getitem__只做轻量的读取和变换。第二异常处理要到位。如果某个样本文件损坏了__getitem__抛异常会导致整个训练崩溃。可以在__getitem__里加 try-except遇到坏样本就返回一个替代样本或者随机跳到另一个索引。但要注意如果坏样本太多这种“跳过”策略会掩盖数据质量问题最好还是先清洗数据。第三注意索引的类型。__getitem__接收的 idx 可能是 Python int也可能是 numpy int64 或 torch tensor。如果你用 idx 去索引一个 Python listnumpy int64 是可以的但 torch tensor 不行。保险的做法是在开头加一句idx int(idx)。4.2 collate_fn 的默认行为与自定义场景默认的collate_fn做的事情很简单把 batch 个样本的 tuple 按位置拆开每个位置上的元素拼成一个 batch。比如每个样本返回(image, label)默认 collate 会返回(batch_images, batch_labels)其中 batch_images 是把所有 image 沿第 0 维堆叠。但很多场景下默认行为不够用。比如目标检测中每张图片的标注框数量不一样没法直接堆叠需要自定义 collate_fn 把标注整理成列表。NLP 中每个句子的长度不一样需要 padding 到同一长度同时生成 attention mask。多模态任务中不同模态的数据需要分别处理可能还要做对齐。写自定义 collate_fn 的时候最容易犯的错误是维度搞混。建议在函数开头打印一下输入的形状确认无误后再写后续逻辑。另外collate_fn 是在主进程里执行的不是在 worker 里所以如果里面有耗时操作会阻塞整个数据流。尽量把重活放在 Dataset 的__getitem__里collate_fn 只做组装。4.3 数据增强与 transform 的放置位置数据增强应该放在 Dataset 的__getitem__里还是放在 collate_fn 里答案是绝大多数情况下放在__getitem__里。原因是__getitem__是在 worker 进程中执行的多个 worker 可以并行做增强充分利用 CPU 多核。如果放在 collate_fn 里就是在主进程串行执行速度会慢很多。但有一个例外如果增强操作需要 batch 内样本间的交互比如 MixUp、CutMix那就必须放在 collate_fn 里因为__getitem__一次只能看到一个样本。另外验证集和测试集的 transform 通常和训练集不一样。训练集用随机裁剪、随机翻转等增强验证集用中心裁剪、不翻转。这个切换可以通过在 Dataset 初始化时传入不同的 transform 来实现也可以通过dataset.transform val_transform动态修改。我一般倾向于前者逻辑更清晰。5. 常见问题与排查技巧实录5.1 多进程相关的典型报错报错一RuntimeError: DataLoader worker (pid xxx) is killed by signal: Bus error这个报错通常是因为共享内存不够。Linux 系统默认的/dev/shm大小一般是内存的一半如果你开了很多 worker每个 worker 又往共享内存里写大量数据就会撑爆。解决办法有两个一是减小 num_workers二是增大/dev/shm的大小需要系统权限。报错二TypeError: cannot pickle _thread.lock object这个报错是因为 Dataset 里包含了不可序列化的对象比如数据库连接、文件句柄、线程锁。在 Windows 上worker 是通过 spawn 创建的需要把 Dataset 序列化后传给子进程所以要求 Dataset 必须可 pickle。解决办法是把这些不可序列化的对象放在__getitem__里懒加载或者用__getstate__和__setstate__自定义序列化行为。报错三RuntimeError: Cowardly refusing to serialize non-leaf tensor which requires_grad这个报错是因为 Dataset 返回了需要梯度的 tensor。DataLoader 在 worker 和主进程之间传递数据时会尝试序列化 tensor而需要梯度的 tensor 不能被序列化。解决办法是在 Dataset 里用torch.no_grad()包裹或者用.detach()去掉梯度。5.2 性能问题的排查思路数据加载慢怎么定位瓶颈我的排查顺序是这样的第一步把 num_workers 设为 0测一下纯数据加载的耗时。如果 num_workers0 时每个 epoch 要 100 秒num_workers8 时要 20 秒说明多进程确实起了作用但还可以继续优化。第二步用time.time()在__getitem__里打点看看时间花在哪里。是文件 IO 慢还是 transform 慢还是其他操作慢。如果是文件 IO 慢考虑把数据转成更高效的格式比如 LMDB、HDF5、WebDataset如果是 transform 慢考虑用 GPU 做增强比如 NVIDIA DALI或者用更快的库比如 OpenCV 替代 PIL。第三步检查 GPU 利用率。用nvidia-smi -l 1或者watch -n 1 nvidia-smi观察。如果 GPU 利用率经常掉到 0说明数据供给跟不上如果 GPU 利用率稳定在 90% 以上说明数据加载不是瓶颈不用再折腾了。第四步检查内存和共享内存使用。用free -h看内存用df -h /dev/shm看共享内存。如果共享内存快满了减小 num_workers 或者减小 prefetch_factor。5.3 常见问题速查表问题现象可能原因解决办法GPU 利用率低训练慢num_workers 太小或为 0增大 num_workers开启 pin_memory内存爆炸Dataset 在__init__里加载了大数据每个 worker 一份副本改为懒加载或用内存映射报 Bus error共享内存不足减小 num_workers 或增大 /dev/shm报 pickle 错误Dataset 含不可序列化对象懒加载或自定义序列化随机增强失效用了 Python 原生 random改用 torch 的随机函数最后一个 batch 报错drop_lastFalse 且 BatchNorm 遇到单样本设 drop_lastTrueshuffle 和 sampler 冲突同时设了多个互斥参数只保留一个Windows 上无限递归没加if __name__ __main__加上入口保护5.4 几个我踩过的坑第一个坑在 Dataset 的__init__里做数据过滤。我曾经写过一个 Dataset在__init__里遍历所有文件把损坏的文件剔除掉。结果 num_workers8 时每个 worker 都做了一遍过滤启动时间翻了 8 倍。后来改成在外部做好过滤把干净的文件列表传进来启动时间立刻降下来了。第二个坑在__getitem__里用random.random()做随机增强。结果发现所有 worker 产生的随机数序列一模一样因为每个 worker 的 Python random 种子相同。改成torch.rand()之后问题解决因为 PyTorch 会为每个 worker 设置不同的种子。第三个坑pin_memoryTrue 但 num_workers0。这种情况下 pin_memory 是在主进程里做的会阻塞训练循环反而更慢。pin_memory 只有在 num_workers 0 时才有意义因为这时数据拷贝可以和主进程的计算重叠。第四个坑用IterableDataset时忘了设batch_sizeNone。IterableDataset是自己控制 batch 的如果 DataLoader 再设 batch_size会报错。正确做法是设batch_sizeNone让 Dataset 自己 yield batch。6. 进阶用法与性能调优6.1 IterableDataset 的适用场景前面讲的都是MapStyleDataset也就是通过索引访问的 Dataset。PyTorch 还提供了IterableDataset它不通过索引访问而是直接 yield 样本。适合以下场景数据流是无限的或者非常大的没法预先知道长度。数据需要按顺序读取比如从数据库游标、日志流、网络流中读取。数据分片存储在多个文件中需要动态决定读取顺序。IterableDataset的多进程处理和MapStyleDataset不一样。因为每个 worker 都会独立地迭代数据如果不做特殊处理多个 worker 会读到重复的数据。解决办法是用worker_init_fn给每个 worker 分配不同的数据分片或者用torch.utils.data.get_worker_info()在运行时获取 worker 编号据此决定读哪些数据。6.2 用 Sampler 实现分布式训练的数据切分分布式训练时每个 GPU 应该只看到数据集的一部分同时保证每个 epoch 所有 GPU 合起来能看到全部数据。DistributedSampler就是干这个的。它的工作原理是根据当前进程的 rank 和总进程数把数据集分成若干份每个进程只取自己那一份。同时每个 epoch 会重新打乱顺序保证不同 epoch 看到的数据组合不同。使用时要注意DistributedSampler需要设置shuffleTrue来打乱同时 DataLoader 的 shuffle 要设为 False。另外每个 epoch 开始前要调用sampler.set_epoch(epoch)否则每个 epoch 的打乱顺序会一样。6.3 数据预取与缓存策略如果数据预处理很耗时可以考虑把处理好的数据缓存起来。最简单的做法是在__getitem__里加一个内存缓存字典第一次读取时做预处理并缓存后续直接返回缓存结果。但要注意缓存大小别把内存撑爆。更高级的做法是用磁盘缓存。比如把预处理后的数据存成 numpy 数组或 tensor 文件下次直接读。或者用 LMDB 这样的嵌入式数据库支持高效的随机读取。还有一种策略是预取。除了 DataLoader 自带的 prefetch_factor你还可以在 Dataset 层面做预取。比如在__getitem__里提前读取下一个样本的数据利用 IO 等待的时间做其他事情。但这会增加代码复杂度一般只在极端性能要求下才用。6.4 性能调优的检查清单调优数据管道时我一般按这个清单逐项检查num_workers 是否设到了合适值观察 GPU 利用率pin_memory 是否开启GPU 训练时prefetch_factor 是否足够观察 GPU 是否有等待Dataset 是否有不必要的重复初始化transform 是否用了高效的库OpenCV 通常比 PIL 快数据格式是否高效避免大量小文件考虑打包格式是否有内存泄漏观察长时间训练的内存变化collate_fn 是否有耗时操作应该尽量轻量7. 一些个人体会DataLoader 这个组件刚入门的时候觉得很简单用久了才发现水深。它的每一个参数背后都有设计者的权衡每一个报错背后都有特定的运行机制。我自己的经验是不要等到出问题了才去查文档而是提前把多进程、内存、序列化这些机制搞清楚后面能省很多时间。另外数据管道的优化是有上限的。当 GPU 利用率稳定在 90% 以上时再优化数据加载的收益就很小了这时候应该把精力放在模型结构、学习率调度、正则化等更影响最终效果的地方。不要陷入“为了优化而优化”的陷阱。最后分享一个习惯每次搭好一个新的数据管道我会先用小数据量跑几个 epoch确认数据读取正确、增强效果符合预期、batch 组装没问题然后再上全量数据。这个习惯帮我避免了很多“训练了半天才发现数据有问题”的尴尬。