训练脚本里有一段代码几乎人人都在用，但真正逐行想过它为什么长这样的人并不多——就是DataLoader(dataset, batch_size=32, shuffle=True)这一行。它藏在每个训练脚本的开头几行，跑起来不报错，看起来也很简单，于是大家复制粘贴过去就再也没动过。可一旦模型训练速度不对、GPU 利用率上不去、loss 曲线抖动得莫名其妙，问题往往就出在这行看起来最无害的代码上。DataLoader 和它背后的 Dataset 抽象，是 PyTorch 数据加载体系里最值得单独抽一讲的部分：它决定了你的数据以什么顺序、什么粒度、什么并发度、什么内存形态送进模型，也决定了你的显卡到底是在算，还是在等。这篇内容面向刚过完张量、自动求导、nn.Module这几关、准备把真实数据集塞进模型的读者，也适合已经能跑通训练但觉得"慢得不正常"的人。我会从手写 batch 循环讲起，把 Dataset 的契约、DataLoader 每个参数的运行时行为、多进程的坑、GPU 拷贝的配合方式，以及一次真实的数据瓶颈排查过程，尽量讲透。

1. 从一段手写取样代码说起：DataLoader 到底替你做了哪几件事

第一次把自建数据塞进模型的人，大概率写过类似这样的东西：拿到一个形状为 (N, D) 的特征矩阵和一个 (N,) 的标签向量，然后用一个for i in range(0, N, batch_size)手动切片，每个 batch 手动转成张量。这段代码能跑，但它同时暴露了四个问题，而 DataLoader 存在的意义就是把这四个问题一次性解决掉。

1.1 没有 DataLoader 时，一个 batch 是怎么被拼出来的

先看下面这段"手动版"训练循环，我把它写全，方便你对照后续的对比：

import numpy as np import torch N, D = 50000, 128 X = np.random.randn(N, D).astype(np.float32) y = np.random.randint(0, 10, size=N).astype(np.int64) batch_size = 64 lr = 1e-3 model = torch.nn.Linear(D, 10) for epoch in range(3): order = np.random.permutation(N) # 手动打乱 for start in range(0, N, batch_size): idx = order[start:start + batch_size] xb = torch.from_numpy(X[idx]) # 手动切片 + 搬运 yb = torch.from_numpy(y[idx]) loss = torch.nn.functional.cross_entropy(model(xb), yb) loss.backward() with torch.no_grad(): for p in model.parameters(): p -= lr * p.grad p.grad = None

这段代码在数据量小的时候毫无问题，可它把四件本来可以解耦的事情揉在了一起：索引采样（permutation加切片）、单条数据构造（这里是 numpy 切片，真实场景可能是打开图片、解码、增强）、组批（把一批样本堆成张量）、设备搬运（from_numpy之后还要再.to(device)）。揉在一起的后果是，你没法单独优化任何一环——想加多进程读取？得自己写multiprocessing。想换成变长输入？得自己写 padding。想按类别加权采样？得自己维护权重表。

DataLoader 的第一层价值就是把"取哪些数据"和"怎么取一条数据"拆开。取哪些数据由 Sampler 负责，怎么取一条数据由 Dataset 负责，组批和搬运由 DataLoader 的收集逻辑负责。这个职责划分不是设计上的洁癖，它直接带来了可替换性：换个采样策略不用动数据集代码，换个数据格式不用动训练循环。

1.2 DataLoader 内部四个环节：采样、取数、组批、搬运

把一次for batch in loader展开，内部大致跑的是这么一条流水线：采样器先产出一批索引（长度等于batch_size），然后拿索引去 Dataset 上逐个调__getitem__拿到单条样本，再交给collate_fn把它们拼成一个张量或字典，接着按需放到页锁定内存里，最后由你手动.to(device)送进显存。理解这条流水线之后，很多参数就不用死记了——它们各自挂在这条流水线的某个环节上。

我在实际项目里见过一个很典型的误解：有人以为把num_workers调大就能加速所有环节。实际上num_workers只作用在"取数"和"组批"这两环，而且具体地说，__getitem__和collate_fn是在子进程里执行的，主进程只负责收结果和搬运。如果瓶颈在采样器（比如采样器本身要读一个大文件算权重），或者瓶颈在 H2D 拷贝（数据量太大而 PCIe 带宽不够），加 worker 是没有用的。这也是为什么排查性能问题必须先定位环节，而不是先调参数。

还有一点值得提前说清楚：DataLoader 本身不负责打开文件。它不知道你读的是 CSV、图片还是音频，全部文件 IO 逻辑都写在 Dataset 的__getitem__里。所以"DataLoader 慢"这个说法在绝大多数情况下是不准确的，慢的是你自己写的__getitem__，DataLoader 只是忠实地按照你设定的并发度去调它。

2. 把 Dataset 写对：三个方法各自的边界在哪

自定义 Dataset 需要实现__init__、__len__、__getitem__三个方法，看起来门槛极低，但我在代码评审里见过的问题，八成集中在这三个方法里该放什么、不该放什么。写对了，后面的 DataLoader 参数才有调优空间；写错了，参数怎么调都是在给一个错误的结构打补丁。

2.1__init__里该放什么、绝对不能放什么

__init__只在 Dataset 对象构造时执行一次，所以它是放"全局一次性准备"的地方：读取索引文件、构建路径列表、加载标签表、初始化不会变的常量。但它有两个禁区。

第一个禁区是不要在__init__里打开需要长期持有的文件句柄或数据库连接。原因在于多进程：Linux 上 DataLoader 默认用 fork 创建 worker，父进程打开的文件描述符会被子进程继承，多个子进程同时读同一个文件偏移量，轻则数据错乱，重则直接抛异常。稳妥的做法是在__getitem__里按需打开、用完关闭，或者用worker_init_fn在子进程内部单独建立连接。

第二个禁区是谨慎把所有数据一次性读进内存。小数据集（几万条、几百 MB）这么做完全没问题，反而能省掉大量的 IO 开销；但如果是几万张高分辨率图片，全读进内存会直接把机器打爆，尤其是开了多个 worker 之后——每个 worker 都有自己的一份 Python 对象，虽然 fork 后有写时复制，但 Python 的引用计数机制会让对象页被频繁标记为脏页，实际内存占用远高于你的预期。

我通常的做法是分三层：路径和标签这类轻量元数据在__init__里全量加载；中等规模的数值特征（比如几千维的向量）在__init__里转成float32的 numpy 数组常驻；大文件（图片、音频、视频）只存路径，在读的时候现取。

import numpy as np import torch from torch.utils.data import Dataset class TabularDataset(Dataset): def __init__(self, features: np.ndarray, labels: np.ndarray): # 一次性转换 dtype，避免每条样本重复做类型转换 self.features = torch.from_numpy(features.astype(np.float32)) self.labels = torch.from_numpy(labels.astype(np.int64)) def __len__(self): return self.features.shape[0] def __getitem__(self, idx): return self.features[idx], self.labels[idx]

这段代码有个细节值得注意：astype(np.float32)放在了__init__里。如果放进__getitem__，那么每条样本每次被取用都要做一次类型转换，几万条样本乘以几十个 epoch，这个开销是实实在在的。同理，如果你有固定的归一化均值方差，也应该在__init__里算好存成张量，而不是每次在__getitem__里重算。

2.2__len__与__getitem__的隐性契约

__len__返回值决定了一个 epoch 有多少个 batch（准确说是ceil(len / batch_size)或floor，取决于drop_last）。它必须是一个稳定的整数，不能随调用次数变化。听起来是废话，但我确实见过在__len__里做过滤逻辑的写法，结果每次调用返回不同值，导致 DataLoader 内部计算 batch 数量错乱，训练循环跑着跑着就提前结束或者多跑一轮。

__getitem__的契约更微妙一点：它接收的索引一定是range(len(self))范围内的整数，返回值的类型则完全自由——张量、numpy 数组、字典、元组都行，只要collate_fn能处理。默认的 collate 逻辑会尝试把一批返回值堆叠成 batch 维度，所以如果你返回的是形状为 (C, H, W) 的单张图片张量，默认 collate 会得到 (B, C, H, W)，这正是大多数人期望的行为。

有一个容易忽略的点：当你用了sampler或者batch_sampler，传给__getitem__的索引未必是顺序的，也未必覆盖全集。比如WeightedRandomSampler会按权重有放回地抽，同一条样本可能在一个 epoch 里出现多次，某些样本可能一次都不出现。所以__getitem__不能对索引顺序做任何假设，更不能在里面维护"我读到了第几条"这类状态。

2.3 多进程下的随机性：为什么你的数据增强像是失效了

这个坑我踩过，而且踩得很难受。现象是：训练集做了随机裁剪和随机翻转，但模型很快就过拟合，验证集指标和训练集指标差距异常小，看起来像"增强没生效"。排查了半天才发现，四个 worker 里每个 epoch 产出的增强模式高度一致，等于每个 batch 里的样本被做了几乎相同的变换。

根因在于随机数种子的继承。在 fork 模式下，每个 worker 都是从父进程 fork 出来的，random和numpy.random的全局状态在 fork 那一刻是相同的，而 PyTorch 只会自动给torch的随机数生成器设置不同的种子（通过base_seed + worker_id）。如果你在__getitem__里用的是random.random()或np.random.rand()，四个 worker 就在用同一串随机数。

修法有两个。最直接的是给 DataLoader 传worker_init_fn：

import random import numpy as np import torch def worker_init_fn(worker_id: int): # torch 已经按 worker_id 设置了不同种子，这里取出来给 numpy 和 random 用 seed = torch.initial_seed() % (2 ** 32) np.random.seed(seed) random.seed(seed) loader = torch.utils.data.DataLoader( dataset, batch_size=64, shuffle=True, num_workers=4, worker_init_fn=worker_init_fn, )

另一个更彻底的做法是让 Dataset 自带一个torch.Generator，在__getitem__里显式传generator=，避免依赖全局状态。后者的好处是可复现性强，缺点是每处随机调用都要记得传参，写起来啰嗦。我现在的习惯是：简单的random/numpy增强用worker_init_fn兜底，需要严格复现的实验就上独立的Generator。

注意：worker_init_fn只在num_workers > 0时才会被调用。当num_workers=0（数据在主进程里取）时，这个函数根本不会执行，你的种子设置也就不会生效。调试阶段如果临时把 worker 关掉，要意识到随机行为会和正式训练时不同。

3. DataLoader 参数逐个拆：每个参数改的是运行时哪一段

DataLoader的构造参数有十几个，文档里逐个罗列看着头大，但按"它挂在流水线哪一环"来分组，记忆量会小很多。我把它们分成三组：决定取哪些数据的、决定取多快的、决定怎么搬运的。这一节按这个分组讲，重点是参数改了之后运行时到底发生了什么变化。

3.1 影响"取哪些数据"的参数：shuffle、sampler、drop_last

shuffle=True本质上是给 DataLoader 装了一个RandomSampler。它在每个 epoch 开始时把索引全量打乱一次，注意是每个 epoch 打乱一次，不是每个 batch 打乱一次。这意味着同一个 epoch 内，一条样本只会被访问一次（除非你用了有放回的采样器）。

sampler和shuffle是互斥的，传了sampler就不能再传shuffle=True，否则直接报错。这个设计是合理的：shuffle只是RandomSampler的语法糖，两者同时存在时哪个优先会非常反直觉。

drop_last决定最后一个不满的 batch 是丢掉还是保留。默认False，也就是保留。这里有个必须知道的坑：如果你的模型里有BatchNorm，而最后一个 batch 恰好只剩 1 条样本，训练模式下的 BatchNorm 会直接报错（大意是每个通道需要多于一个值来计算方差）。数据集大小不能被batch_size整除时，这个情况一定会出现。所以只要用了 BatchNorm，我基本都会顺手加drop_last=True，代价是每个 epoch 丢掉不到一个 batch 的数据，换来的是稳定性。

loader = torch.utils.data.DataLoader( dataset, batch_size=64, shuffle=True, drop_last=True, # 配合 BatchNorm 使用，避免尾批只有 1 条样本 num_workers=4, )

如果你的任务不能丢数据（比如验证集必须全量评估），那就把验证集的 DataLoader 设成drop_last=False，同时确保模型在评估模式下跑（model.eval()），评估模式下的 BatchNorm 用的是滑动统计量，不会因为批大小为 1 报错。

3.2 影响"取多快"的参数：num_workers、prefetch_factor、persistent_workers

num_workers是最常被调、也最常被误解的参数。它指定子进程数量，0表示不开子进程、全在主进程里做。开了子进程之后，主进程可以一边等 GPU 算上一个 batch，一边让子进程准备下一个 batch，这是流水线并行的基础。

prefetch_factor控制每个 worker 提前准备多少个 batch，默认是 2（仅在num_workers > 0时生效）。把它调大相当于加长了缓冲区，对"取数耗时抖动大"的场景有帮助（比如有的图片大有的小），代价是内存占用上升。我一般不去动它，除非确认存在明显的取数抖动。

persistent_workers是我强烈建议开的参数（PyTorch 1.7+）。默认情况下，每个 epoch 结束时 worker 进程会被销毁，下一个 epoch 开始时重新创建。如果你的数据集对象在__init__里做了重活（比如加载了一个索引表），这个重建成本会每个 epoch 付一次。设成True之后，worker 进程会一直存活到 loader 被删除，省掉这部分开销。小数据集上差异不明显，几万条以上、num_workers又开得比较多的时候，能省下不少时间。

3.3 影响"怎么搬"的参数：pin_memory、collate_fn、batch_size=None

pin_memory=True让 DataLoader 把组好的 batch 放进页锁定内存（pinned memory），这样从 CPU 到 GPU 的拷贝可以走 DMA，并且能配合non_blocking=True做异步传输。它的作用点在"搬运"环节，不改变取数速度。下一节会细讲。

collate_fn是自定义组批逻辑的入口，默认实现是default_collate，它能处理张量、numpy 数组、数字、字符串、列表、字典的嵌套组合。当你的样本是变长序列、或者需要返回额外的元信息（比如原始尺寸、样本 ID），就得自己写。有一个使用细节：collate_fn是在 worker 子进程里执行的，所以它里面不要做需要访问主进程状态的事情，也不要往里面塞重活。

batch_size=None是个很少用但很特殊的取值，它表示"不做自动组批"，此时 Dataset 返回什么，DataLoader 就原样吐什么，迭代器产出的就是单条样本。配合batch_sampler使用时，相当于把组批的控制权完全交给你。我在做序列标注这类"一个 batch 内的样本必须来自同一个长文档"的场景时用过这种方式，用自定义的batch_sampler按文档分组。

4. num_workers 调到几才算合适：实测曲线与多进程的几类硬伤

关于num_workers设多少，网上的建议从 2 到 CPU 核数都有，这些数字单独看没有意义，因为它取决于你的__getitem__有多重、磁盘有多快、CPU 核数有多少、有没有和别的进程抢资源。我给不出一个通用数字，但可以给你一条实测曲线和一套判断方法。

4.1 收益不是线性的：从 0 到 16 的实测数据

下面这组数据来自我自己搭的一个测试场景：约五万张中等分辨率的图片，__getitem__里做解码、短边缩放到 256、随机裁剪到 224、转张量归一化，机械硬盘存储，8 核 CPU，单卡训练。只测数据吞吐，不含模型前向反向：

曲线在 8 附近就基本平了，再往上加几乎没有收益，反而因为进程调度和内存带宽竞争开始轻微下降。这是很典型的形态：加速比永远达不到 worker 数量的线性倍数，因为主进程收结果、collate_fn、以及磁盘 IO 带宽都会成为新的瓶颈。8 核机器上 8 个 worker 已经是把 CPU 吃满的状态了，再加就是自己和自己抢。

由此得出一个我一直在用的经验法则：num_workers不要超过分配给这个训练任务的物理核数，多卡训练时更要注意，四张卡各开 8 个 worker 就是 32 个进程在抢同一批核，实际效果往往不如每卡 4 个。另外，如果__getitem__很轻（比如纯内存里的张量切片），num_workers=0甚至可能比开进程更快，因为省掉了进程间传数据的序列化开销——小张量的 IPC 成本占比会很高。

4.2 跨平台差异：Windows 的 spawn 与 Linux 的 fork 各自埋的雷

多进程的启动方式在两个平台上完全不同，这是很多人"本地能跑、服务器报错"或者反过来的根源。

Linux 默认fork，子进程直接继承父进程的内存状态，启动快、开销小。它的雷在于继承：父进程里打开的数据库连接、日志文件句柄、全局锁，都会被复制到子进程，多个子进程同时用同一个连接，出问题是必然的。我在一个项目里遇到过父进程初始化了一个 sqlite 连接放在全局变量里，num_workers=0时一切正常，开了 4 个 worker 之后随机报"database is locked"，查了很久才定位到 fork 继承。

Windows 和 macOS 默认spawn，子进程是全新启动的 Python 解释器，需要重新 import 主模块。它的雷是必须把训练入口包在if __name__ == "__main__":里，否则子进程在 import 主模块时又会执行一遍创建 DataLoader 的代码，形成无限递归，直接抛RuntimeError: An attempt has been made to start a new process before the current process has finished its bootstrapping phase。

import torch from torch.utils.data import DataLoader def build_loader(dataset): return DataLoader(dataset, batch_size=64, shuffle=True, num_workers=4) if __name__ == "__main__": # spawn 模式下，这段必须放在 main guard 里 dataset = MyDataset("data/train") loader = build_loader(dataset) for epoch in range(10): for batch in loader: pass

即使你只在 Linux 上跑，我也建议养成写main guard的习惯，因为代码迟早会在别人的机器上跑，或者被torch.multiprocessing的其他组件间接调用。

4.3 worker 被杀、内存暴涨、句柄失效的排查方向

num_workers相关的报错信息通常不太友好，我把见过的几类整理成一张对照表：

最后那条排查技巧特别有用：遇到 worker 相关的诡异报错，第一件事就是把num_workers调成 0 复现。主进程模式下异常堆栈是完整的，能直接指向__getitem__的某一行；而在子进程里，你往往只能看到一个干巴巴的进程退出信息。等定位到具体代码再改回多进程验证，这个顺序能省下大量时间。

5. 让 GPU 不再等数据：pin_memory 与异步拷贝的配合姿势

到这里，数据已经能被稳定、快速地取出来了，但它还在 CPU 内存里。从 CPU 内存到显存这一段，是另一个独立的耗时环节，也是很多人漏掉的一环。我见过不少训练脚本，数据加载已经调得不错，GPU 利用率却还是呈现规律的锯齿状，问题就出在拷贝没有异步化。

5.1 一次迭代的时间到底花在哪

把一个训练迭代拆成四段打点，通常能看到这样的分布（数值是某次图像分类任务的实际测量，仅供理解比例）：

如果拷贝是同步的，那么每个迭代实际耗时是 42 + 3.2 = 45.2 ms，其中 3.2 ms 里 GPU 完全闲着。单看比例不高（7%），但乘以几万个迭代就是可观的时间。如果数据量更大（比如高分辨率图像或长序列），拷贝耗时会涨到十几毫秒甚至更多，这时候 GPU 的等待就会变得非常明显。

5.2 pin_memory + non_blocking 的最小改动

打开异步拷贝只需要改两处，没有额外依赖：

loader = torch.utils.data.DataLoader( dataset, batch_size=64, shuffle=True, num_workers=8, pin_memory=True, # 组好的 batch 放进页锁定内存 persistent_workers=True, ) device = torch.device("cuda") for xb, yb in loader: # non_blocking=True 让拷贝和后续计算在不同流上重叠 xb = xb.to(device, non_blocking=True) yb = yb.to(device, non_blocking=True) ...

这里的关键机制是：页锁定内存不会被操作系统换出到磁盘，GPU 的 DMA 引擎可以直接从这块内存读取数据，不需要 CPU 参与逐字节搬运。而non_blocking=True告诉 PyTorch 把这次拷贝提交到拷贝流之后就立刻返回，CPU 可以继续往下走。两者必须配合使用，只开pin_memory不开non_blocking，收益非常有限；只开non_blocking不开pin_memory，拷贝实际上仍然是同步的（因为源内存是可分页的），甚至会因为要走暂存缓冲区而略慢。

有一个使用上的注意点：non_blocking=True之后，拷贝完成和计算开始之间没有显式的同步点，PyTorch 靠流依赖保证顺序，一般不需要你手动synchronize()。但如果你在拷贝之后立刻在 CPU 侧读这个张量的值（比如打印、item()、或者传给 numpy），就会触发隐式同步，反而可能更慢。所以调试打印尽量放在loss.item()这种必须同步的位置，不要散布在各个中间步骤。

5.3 什么时候不该开 pin_memory

pin_memory不是无条件的好。页锁定内存是有限资源，操作系统不能把它换出，开得太多会影响整机其他进程的内存分配。以下几种情况我会关掉它：

CPU 上训练时完全不需要，因为没有 H2D 拷贝。数据集特别小、单批数据只有几百 KB 时，页锁定的管理开销可能超过收益，实测差别在噪声范围内。机器内存本来就紧张（比如同时跑多个训练任务）时，大量页锁定内存会挤压其他进程，严重时触发整机内存告警。还有一种情况是数据在 GPU 上直接生成（比如某些合成数据任务），压根没有 CPU 到 GPU 的搬运，这时候pin_memory纯属浪费。

提示：pin_memory_device参数只在较新版本里才有，用于指定锁页内存对应的目标设备，多卡场景下如果发现拷贝没有被正确路由，可以查一下这个参数的版本支持情况，别照着新文档在旧版本上硬套。

6. 不均衡样本、变长序列：采样器与 collate_fn 的组合用法

前面几节处理的是"通用场景"，但真实项目里总会遇到两类特殊数据：一类是类别严重倾斜，一类是长度参差不齐。这两种情况用默认配置跑，要么模型学偏，要么显存被 padding 撑爆，都得靠采样器和collate_fn来解决。

6.1 WeightedRandomSampler 处理类别倾斜

假设一个二分类任务，正负样本比例是 1:20。直接shuffle=True训练，模型很快会学会"全预测为多数类"，准确率看着有 95%，但少数类的召回率接近 0。最省事的修法是按类别频率的倒数给每条样本赋权，让采样器在抽 batch 时更倾向于抽少数类。

import torch from torch.utils.data import DataLoader, WeightedRandomSampler # labels 是全量标签，形状 (N,) class_counts = torch.bincount(labels) class_weights = 1.0 / class_counts.float() sample_weights = class_weights[labels] # 每条样本的权重 sampler = WeightedRandomSampler( weights=sample_weights, num_samples=len(sample_weights), # 一个 epoch 抽这么多次 replacement=True, # 有放回，否则少数类抽完就没了 ) loader = DataLoader(dataset, batch_size=64, sampler=sampler, num_workers=4)

两个细节值得说清楚。replacement=True是必须的，因为权重不均等时，如果不放回，采样器在抽完所有少数类之后就没得抽了，行为会变得很奇怪。num_samples设成数据集大小，是为了让一个 epoch 的迭代次数和原来保持一致，方便对比实验；设得更大相当于延长了 epoch，会改变学习率调度的时间尺度，容易把实验搞乱。

另外要注意，用了WeightedRandomSampler之后，一个 epoch 里被抽到的样本集合和全集不是一回事，某些样本会被重复抽多次，某些一次都不出现。所以用它算出来的训练集指标和用普通 loader 算出来的不完全可比，评估一定要用独立的、shuffle=False的验证集 loader。

6.2 collate_fn 处理变长输入，顺便把 padding 的浪费降下来

变长序列的默认collate_fn会直接报错，因为张量形状不一致没法堆叠。标准解法是写一个按 batch 内最大长度做 padding 的 collate：

import torch def pad_collate(batch): seqs, labels = zip(*batch) # batch 是 [(seq, label), ...] lengths = torch.tensor([len(s) for s in seqs], dtype=torch.long) max_len = int(lengths.max()) padded = torch.zeros(len(seqs), max_len, dtype=torch.long) for i, s in enumerate(seqs): padded[i, :len(s)] = s return padded, lengths, torch.tensor(labels) loader = DataLoader(dataset, batch_size=32, collate_fn=pad_collate, num_workers=4)

这里返回了lengths，目的是让模型能拿到真实长度去做 mask，避免 padding 位置参与注意力计算或损失计算。如果不返回长度，后面还得靠pad_id反推，既慢又容易出错。

padding 到 batch 内最大长度这件事本身是有浪费的：如果一个 batch 里大部分序列长度是 20，偶尔混进来一条 500，那整个 batch 都要 padding 到 500，计算量凭空多了二十几倍。缓解办法是按长度分桶——先按长度排序，再把相邻的若干条组成一个 batch，这样同一个 batch 内长度接近，padding 浪费就小了。实现上可以写一个按长度分组的batch_sampler，或者更简单一点：在 Dataset 的__init__里按长度排序，然后用shuffle之外的自定义采样器在每个 epoch 内部做局部打乱。前者更规范，后者更快出效果，看你的时间预算。

6.3 DistributedSampler 的 set_epoch 不能忘

多卡训练时，每个进程必须只看到数据的一个不重叠子集，否则等于变相放大了 batch 又引入了重复样本。DistributedSampler就是干这个的，它会按进程数切分索引。用法上有一个必须做但极容易忘的动作：每个 epoch 开始时调用sampler.set_epoch(epoch)。

from torch.utils.data.distributed import DistributedSampler sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True) loader = DataLoader(dataset, batch_size=64, sampler=sampler, num_workers=4) for epoch in range(num_epochs): sampler.set_epoch(epoch) # 忘了这行，每个 epoch 的切分顺序完全一样 for batch in loader: ...

不调用set_epoch的后果是：每个 epoch 的分片方式完全相同，相当于数据顺序固定，shuffle形同虚设。这个 bug 不会报错，只会让你的训练效果比预期差一点，非常隐蔽。我在别人的代码里见过不止一次，加一行就好了。

7. 一次数据瓶颈的完整定位过程：从 GPU 利用率 30% 到吃满

前面讲的都是分解动作，这一节把它们串起来，还原一次真实的排查过程。我把当时的现象、每一步的判断依据、以及最终的改动都记下来，你可以把这套流程套到自己的项目上。

7.1 先证明瓶颈在数据侧，而不是模型侧

现象是：单卡训练一个图像分类模型，nvidia-smi显示 GPU 利用率在 10% 到 60% 之间反复横跳，平均下来 30% 出头，一个 epoch 要跑四十多分钟，比预估的慢了将近两倍。

第一步不是改代码，而是确认到底是谁在等谁。最省事的判断方法是写一个纯数据迭代的脚本，把模型完全摘掉，只跑 DataLoader，统计吞吐：

import time def benchmark(loader, warmup=3): it = iter(loader) for _ in range(warmup): # 预热，避开进程启动和缓存冷启动 next(it) t0 = time.time() n = 0 for batch in loader: n += batch[0].shape[0] dt = time.time() - t0 return n / dt print(f"{benchmark(loader):.1f} samples/s")

测出来的数字是每秒 210 条左右。而模型侧的理论吞吐（用单 batch 前向反向耗时反推）大约是每秒 520 条。数据供给只有需求的 40%，瓶颈基本可以确定在数据侧，模型本身没问题。这个判断很关键，因为如果反过来是模型侧慢，那调 DataLoader 参数再怎么调都是在做无用功。

7.2 逐层打点，把耗时拆到毫秒

确认瓶颈在数据侧之后，下一步是定位到具体环节。我用了两个层次的打点。

第一个层次是对比num_workers的影响，快速判断是不是并发度不够：0的时候 105 条/秒，4的时候 185 条/秒，8的时候 210 条/秒，16的时候 205 条/秒。曲线在 8 就平了，说明单纯加进程已经到顶，存在别的瓶颈。

第二个层次是在__getitem__内部手动计时，把耗时拆到子步骤：

import time class TimedImageDataset(Dataset): def __init__(self, paths): self.paths = paths self.t_decode = 0.0 self.t_resize = 0.0 self.t_aug = 0.0 def __getitem__(self, idx): t0 = time.perf_counter() img = self._decode(self.paths[idx]) t1 = time.perf_counter() img = self._resize(img) t2 = time.perf_counter() img = self._augment(img) t3 = time.perf_counter() self.t_decode += t1 - t0 self.t_resize += t2 - t1 self.t_aug += t3 - t2 return img

提示：这些累加变量只在num_workers=0时有意义。开了多进程之后，子进程里的self是副本，累加结果不会传回主进程，你会看到全是 0。所以计时调试阶段要先把 worker 关掉。

打点结果显示：解码占了 62%，缩放占了 21%，增强只占 9%，剩下是杂项。解码是大头，而且原始图片是几 MB 一张的相机直出图，每次都要完整解码再缩放，等于绝大部分像素解出来就被丢掉了。

7.3 优化落地与前后对比

定位到问题后，改动方向就很明确了，一共三步。

第一，把缩放提前到存储阶段。离线和一次性地把原图统一缩到短边 320 左右存成质量 90 的 JPEG，训练时只做小幅度缩放和裁剪。这样解码的像素量从千万级降到十万级，解码耗时直接降了一个数量级。这一步的代价是磁盘多占一份（约 30%），收益是所有后续实验都受益，属于一次性投入。

第二，调换__getitem__内的操作顺序。原来是先转换成浮点张量再做随机裁剪，改成先在整型、小尺寸的图像上裁剪，最后才转浮点张量并归一化。浮点转换和归一化涉及整张图的逐元素运算，在裁剪之后做，计算量按裁剪面积等比下降。

第三，调整num_workers与persistent_workers。在解码变快之后，瓶颈转移到了别处，我重新测了一遍曲线，最终定在num_workers=6、persistent_workers=True、prefetch_factor=4，让取数抖动被缓冲区吸收掉。

改动前后的对比：

优化之后 GPU 利用率仍然不是 100%，这很正常——每个迭代末尾总有一点等待，锯齿形是数据管线的固有特征，只要平均利用率上去了、且锯齿不深，就说明供给跟得上。追求 100% 通常意味着数据准备过度超前，反而是浪费了 CPU 资源。

最后分享一个我后来养成的小习惯：每次新开一个训练任务，先花十分钟跑一遍纯数据的 benchmark，把吞吐数字记在实验笔记里，再对照模型的理论吞吐看是否匹配。这个动作花不了多少时间，但能在实验跑到一半发现"怎么这么慢"之前，就把问题挡在门外。数据管线的调优不是一次性的工作，换数据集、换输入尺寸、换增强策略，瓶颈都可能转移，定期重测比一次性调好更实际。