PyTorch DataLoader性能调优：从GPU利用率低到高效数据供给

1. 从一次训练日志异常说起：GPU利用率为什么上不去

如果你跑过PyTorch训练任务，大概率见过这样的场景：nvidia-smi里GPU利用率在15%到40%之间反复横跳，偶尔冲到80%又迅速掉下来，训练一个epoch的时间比预期多了两三倍。你检查了模型结构，参数量不算大；检查了显卡型号，显存也够用；甚至换了更大的batch size，情况依然没有明显改善。这时候问题大概率不在模型本身，而在数据供给这条链路上——也就是DataLoader。

DataLoader是PyTorch里负责把数据集样本组装成batch、再喂给模型的组件。它看起来只是一个简单的迭代器，但背后涉及Dataset.__getitem__的调用、collate_fn的拼装、多进程worker的调度、共享内存的传输、以及CPU到GPU的拷贝。任何一个环节出现瓶颈，GPU就会陷入“等数据”的状态。GPU利用率低，本质上是GPU在空转等CPU把数据准备好。

这篇内容面向的是已经能跑通PyTorch训练、但对性能排查还没有系统方法的开发者。我会从实际排查链路出发，讲清楚怎么定位瓶颈到底在磁盘IO、在CPU预处理、在worker数量配置、还是在数据传输方式上，并给出可以直接复现的验证代码和调优手段。全文围绕PyTorch、DataLoader、GPU利用率、性能排查这几个核心关键词展开，不堆理论，只讲能落地的操作。

需要先明确一个判断基准：GPU利用率低不一定是数据加载的问题。如果模型本身有大量小算子、频繁的CPU-GPU同步、或者用了效率很低的attention实现，GPU也会利用率低。所以在动手调DataLoader之前，先用一个简单方法确认瓶颈方向——把数据加载部分替换成纯内存的随机张量，如果GPU利用率立刻上去了，那问题就在数据侧；如果还是低，那要去看模型和训练循环。这个判断步骤后面会详细展开。

2. 先定位再动手：判断瓶颈是否真的在DataLoader

2.1 用“假数据”做对照实验

排查性能问题最忌讳一上来就改参数。我习惯的做法是先做一个对照实验：构造一个完全不走磁盘、不走Dataset的DataLoader，直接返回随机生成的张量，然后跑几十个step，观察GPU利用率。

import torch from torch.utils.data import DataLoader, TensorDataset # 构造纯内存数据集，模拟“理想情况”下的数据供给 images = torch.randn(10000, 3, 224, 224) labels = torch.randint(0, 1000, (10000,)) dataset = TensorDataset(images, labels) loader = DataLoader(dataset, batch_size=64, shuffle=True, num_workers=0) model = torch.nn.Linear(3 * 224 * 224, 1000).cuda() optimizer = torch.optim.SGD(model.parameters(), lr=0.01) for i, (x, y) in enumerate(loader): x, y = x.cuda(), y.cuda() x = x.view(x.size(0), -1) loss = torch.nn.functional.cross_entropy(model(x), y) loss.backward() optimizer.step() optimizer.zero_grad() if i == 50: break

跑这段代码的同时开另一个终端执行watch -n 0.5 nvidia-smi。如果GPU利用率能稳定在70%以上，说明模型和训练循环本身没问题，瓶颈在真实的数据加载链路。如果这段代码GPU利用率也很低，那要先去排查模型前向反向的计算密度、是否有频繁的.item()调用、是否有不必要的CPU-GPU同步。

这个对照实验的价值在于，它把“数据加载”这个变量从整个训练流程里剥离出来了。很多人跳过这一步，直接去调num_workers，结果调了半天发现根本不是worker的问题，白白浪费时间。

2.2 用PyTorch Profiler看时间都花在哪

对照实验确认瓶颈在数据侧之后，下一步是精确定位时间消耗。PyTorch自带的torch.profiler可以给出每个算子、每个CPU线程的耗时分布。

from torch.profiler import profile, ProfilerActivity with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], schedule=torch.profiler.schedule(wait=1, warmup=1, active=3), on_trace_ready=torch.profiler.tensorboard_trace_handler('./log')) as prof: for step, (x, y) in enumerate(loader): x, y = x.cuda(), y.cuda() loss = model(x) loss.backward() optimizer.step() optimizer.zero_grad() prof.step() if step >= 5: break

跑完之后用TensorBoard打开./log目录，重点看两个东西：一是DataLoader相关的CPU时间占比，二是cudaMemcpyAsync（H2D拷贝）的耗时。如果__getitem__或者collate_fn占了大量CPU时间，说明预处理是瓶颈；如果H2D拷贝时间很长，说明数据传输方式有问题；如果worker进程处于空闲等待状态，说明worker数量或者数据分片策略不合理。

我实际排查过一个图像分类任务，profiler显示每个step有将近60%的时间花在PIL.Image.open和transforms.Resize上，GPU利用率只有20%出头。这就是典型的CPU预处理瓶颈，后面会讲怎么用DataLoader的多进程和更高效的解码库来解决。

2.3 区分“CPU瓶颈”和“IO瓶颈”的快速方法

CPU预处理慢和磁盘IO慢，表现都是GPU等数据，但解决思路完全不同。区分方法很简单：把数据集缓存到内存或者tmpfs里再跑一遍。如果速度明显提升，说明是IO瓶颈；如果没变化，说明是CPU计算瓶颈。

# 把数据拷贝到内存文件系统（Linux） mkdir /dev/shm/dataset_cache cp -r /path/to/dataset/* /dev/shm/dataset_cache/

然后把Dataset的根目录指向/dev/shm/dataset_cache再跑一次。这个操作在Linux上很实用，/dev/shm是内存映射的临时文件系统，读写速度远超普通磁盘。如果数据集体量不大（比如几十GB以内），直接把数据放进去跑，能排除掉IO因素的干扰。

注意：/dev/shm默认大小通常是物理内存的一半，数据量大的话要先确认空间够不够，用df -h /dev/shm查看。

3. num_workers不是越大越好：多进程加载的配置逻辑

3.1 worker数量与CPU核心数的关系

num_workers是DataLoader最常被调整的参数，但很多人要么设成0（单进程），要么无脑设成8或16。实际上worker数量的合理值取决于三个因素：CPU物理核心数、每个样本的预处理耗时、以及内存带宽。

一个经验公式是：num_workers = CPU物理核心数 / (1 + 单样本预处理时间 / 单样本GPU计算时间)。但这个公式在实际中很难精确计算，更实用的做法是从num_workers=4开始，以2为步长往上加，观察GPU利用率和每秒处理的样本数（samples/sec），找到拐点。

import time for nw in [0, 2, 4, 8, 12, 16]: loader = DataLoader(dataset, batch_size=64, shuffle=True, num_workers=nw, pin_memory=True) start = time.time() count = 0 for x, y in loader: x, y = x.cuda(), y.cuda() # 这里放你的模型前向反向 count += x.size(0) if count >= 6400: # 跑够一定样本数就停 break elapsed = time.time() - start print(f"num_workers={nw}, samples/sec={count/elapsed:.1f}")

这段代码跑下来，通常会看到samples/sec先上升后趋于平缓，甚至在某些点上下降。下降的原因一般是worker进程之间的上下文切换开销、共享内存竞争、或者CPU核心被占满导致主进程调度延迟。

3.2 worker数量过多的副作用

worker不是越多越好，这一点我在多个项目里反复验证过。当num_workers超过CPU物理核心数时，会出现几个问题：

第一，worker进程之间争抢CPU时间片，导致每个worker的预处理速度都变慢，整体吞吐反而下降。第二，每个worker都会复制一份数据集对象的引用，如果Dataset里持有大量内存数据（比如把整个数据集load进了内存），内存占用会成倍增长。第三，worker启动和销毁本身有开销，如果每个epoch数据量不大，频繁重建worker反而拖慢训练。

还有一个容易被忽略的点：num_workers > 0时，主进程和worker之间通过共享内存传输数据。如果batch里的张量很大（比如高分辨率图像），共享内存的拷贝开销会变得显著。这种情况下，适当减小batch size或者用更紧凑的数据类型（如float16）可能比增加worker更有效。

3.3 persistent_workers与prefetch_factor的配合

PyTorch 1.7之后引入了persistent_workers参数。默认情况下，每个epoch结束后worker进程会被销毁，下一个epoch重新创建。如果num_workers较大且epoch数很多，这个反复创建销毁的开销不可忽视。设置persistent_workers=True可以让worker在epoch之间保持存活。

loader = DataLoader(dataset, batch_size=64, shuffle=True, num_workers=8, persistent_workers=True, prefetch_factor=4, pin_memory=True)

prefetch_factor控制每个worker预先取多少个batch。默认值是2，意味着每个worker会提前准备2个batch的数据。在GPU计算时间较长、数据加载相对较慢的场景下，适当增大prefetch_factor可以让worker更早开始准备后续数据，减少GPU等待。但设太大也会增加内存占用，一般设2到4比较稳妥。

提示：persistent_workers=True必须和num_workers > 0一起使用，否则会报错。另外，如果数据集在每个epoch会动态变化（比如用了自定义的sampler改变数据顺序），要确认worker保持存活不会导致数据状态错乱。

4. Dataset与预处理：把CPU从繁重的解码任务里解放出来

4.1 __getitem__里的隐形耗时

Dataset.__getitem__是数据加载链路的起点，也是最能藏性能问题的地方。很多教程里的示例代码长这样：

class MyDataset(Dataset): def __getitem__(self, idx): img = Image.open(self.paths[idx]).convert('RGB') img = self.transform(img) return img, self.labels[idx]

这段代码在单进程下跑小数据集没问题，但在多worker场景下，每个worker都要独立执行Image.open和convert。PIL的解码速度在CPU上并不快，尤其是JPEG格式，一张1080p图片解码可能要几毫秒到十几毫秒。如果batch size是64，每个batch光解码就要几百毫秒，GPU自然等不及。

优化方向有几个：一是换用更快的解码库，比如turbojpeg或者opencv-python的cv2.imdecode，后者在多数场景下比PIL快2到3倍。二是把解码后的数据提前缓存，避免每个epoch重复解码。三是用DALI这类GPU加速的数据加载库，把解码和增强都放到GPU上做。

4.2 用cv2替代PIL的实测对比

我做过一个简单的对比测试，在同一台机器上解码1000张1280x720的JPEG图片：

cv2.imdecode的用法要注意，它读出来是BGR格式，需要转成RGB：

import cv2 import numpy as np def read_image_cv2(path): img = cv2.imread(path, cv2.IMREAD_COLOR) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) return img

如果数据集是JPEG且对解码速度要求极高，turbojpeg是更好的选择，但它的安装稍微麻烦一些，需要编译libjpeg-turbo。对于大多数场景，cv2已经能带来明显的提升。

4.3 把增强操作从CPU搬到GPU

数据增强是另一个CPU大户。torchvision.transforms里的RandomResizedCrop、ColorJitter、RandomRotation等操作都在CPU上执行，而且很多是逐像素的Python循环，速度很慢。一个有效的策略是把增强操作尽量简化，或者用torchvision.transforms.v2里的张量版本，它们对批量数据的处理效率更高。

更激进的做法是把增强放到GPU上。比如用kornia库，它提供了GPU版本的图像增强算子：

import kornia.augmentation as K aug = K.AugmentationSequential( K.RandomHorizontalFlip(p=0.5), K.RandomResizedCrop(size=(224, 224)), K.ColorJitter(0.2, 0.2, 0.2, 0.1, p=0.5), data_keys=["input"], ) # 在训练循环里，数据搬到GPU之后再增强 x = x.cuda() x = aug(x)

这样做的好处是CPU只需要负责解码和最基本的张量转换，增强的计算压力转移到了GPU。但要注意，GPU增强会占用一部分GPU算力，如果模型本身已经吃满了GPU，这样做反而可能拖慢训练。适合的场景是模型计算量不大、GPU利用率低、CPU增强成为瓶颈的情况。

4.4 数据预取与缓存策略

如果数据集不大，最直接的办法是在Dataset.__init__里把所有数据加载到内存。这样__getitem__只是做索引和轻量变换，速度极快。

class CachedDataset(Dataset): def __init__(self, paths, labels, transform=None): self.images = [read_image_cv2(p) for p in paths] self.labels = labels self.transform = transform def __getitem__(self, idx): img = self.images[idx] if self.transform: img = self.transform(img) return img, self.labels[idx]

这种方式的代价是内存占用。假设数据集有10万张224x224的RGB图片，每张占224*224*3字节约150KB，总共约15GB。如果内存够大，这是最省事的方案。内存不够的话，可以考虑只缓存解码后的张量到磁盘（比如用numpy.memmap或者lmdb），下次读取时跳过解码步骤。

5. pin_memory与数据传输：CPU到GPU拷贝的优化空间

5.1 pin_memory到底做了什么

pin_memory=True是DataLoader里另一个常被提及的参数。它的作用是把CPU内存中的张量分配到“锁页内存”（pinned memory）里。普通内存可以被操作系统换出到磁盘，而锁页内存不会被换出，因此GPU可以直接通过DMA（直接内存访问）从锁页内存读取数据，不需要CPU介入拷贝。

开启pin_memory之后，x.cuda()这个操作的耗时会明显降低。实测在一个图像分类任务里，pin_memory=False时H2D拷贝占每个step时间的15%左右，开启后降到5%以下。

loader = DataLoader(dataset, batch_size=64, shuffle=True, num_workers=8, pin_memory=True)

但pin_memory不是没有代价的。锁页内存的分配和释放比普通内存慢，而且会占用更多的物理内存。如果数据集很大、batch很多，锁页内存的占用会累积。另外，pin_memory只在num_workers > 0时有明显效果，单进程模式下收益有限。

5.2 non_blocking拷贝的正确用法

配合pin_memory，在把数据搬到GPU时应该使用non_blocking=True：

x, y = x.cuda(non_blocking=True), y.cuda(non_blocking=True)

non_blocking=True表示拷贝操作是异步的，CPU可以在拷贝进行的同时继续执行后续代码。但这里有一个关键前提：只有从锁页内存发起的拷贝才能真正异步。如果pin_memory=False，non_blocking=True实际上还是同步的，甚至可能因为额外的检查而略微变慢。

还有一个容易踩的坑：在non_blocking=True的情况下，如果紧接着对张量做原地操作或者读取它的值，可能会读到未完成拷贝的数据。PyTorch在大多数情况下会自动处理同步，但在自定义的CUDA kernel或者手动管理流的情况下要格外小心。

5.3 用更大的batch减少拷贝次数

H2D拷贝的开销和拷贝次数成正比，和单次拷贝的数据量关系不大（在一定范围内）。也就是说，拷贝一个64样本的batch和拷贝一个256样本的batch，耗时差距远小于4倍。因此，在显存允许的前提下，增大batch size可以有效摊薄拷贝开销。

但batch size增大会影响收敛性，可能需要调整学习率。另外，如果数据加载本身是瓶颈，增大batch size只是让GPU等得更久，并不能解决根本问题。所以这个手段要配合前面的worker和预处理优化一起用。

6. 一套可复用的排查流程与参数模板

6.1 从现象到根因的排查链路

把前面的内容串起来，形成一套可复用的排查流程：

第一步，用假数据对照实验确认瓶颈是否在数据侧。如果假数据下GPU利用率正常，继续；否则去查模型和训练循环。

第二步，用torch.profiler定位时间消耗。重点看__getitem__、collate_fn、H2D拷贝三部分的占比。

第三步，用/dev/shm缓存排除IO因素。如果缓存后速度提升明显，考虑用更快的存储或者数据预取。

第四步，调整num_workers，从4开始以2为步长测试，找到samples/sec的拐点。

第五步，优化__getitem__里的解码和增强操作，换用cv2或turbojpeg，考虑GPU增强。

第六步，开启pin_memory和non_blocking，确认拷贝开销降到最低。

第七步，如果以上都做了还是不够，考虑persistent_workers、prefetch_factor，或者引入DALI等专用数据加载库。

6.2 一份可以直接抄的DataLoader配置

综合以上经验，我给出一份适用于大多数图像分类任务的DataLoader配置模板：

from torch.utils.data import DataLoader loader = DataLoader( dataset, batch_size=64, # 根据显存调整 shuffle=True, num_workers=8, # 从CPU物理核心数的一半开始试 pin_memory=True, # 开启锁页内存 persistent_workers=True, # 避免epoch间重建worker prefetch_factor=4, # 每个worker预取4个batch drop_last=True, # 避免最后一个不完整batch ) # 训练循环里 for x, y in loader: x = x.cuda(non_blocking=True) y = y.cuda(non_blocking=True) # ... 模型前向反向

这份配置不是万能的，但作为一个起点，它能覆盖大部分中等规模图像任务的场景。实际使用时根据profiler的结果微调num_workers和prefetch_factor。

6.3 几个容易忽略的细节

第一，collate_fn的默认实现是torch.stack，如果batch里的样本形状不一致（比如变长序列），需要自定义collate_fn。自定义时要注意不要在collate_fn里做耗时的CPU计算，它是在主进程里执行的，会阻塞数据供给。

第二，如果用了IterableDataset，num_workers的行为和MapStyleDataset不同。每个worker会独立遍历数据流，需要手动用worker_init_fn做数据分片，否则多个worker会读到重复数据。

第三，Windows上num_workers > 0需要把训练代码放在if __name__ == '__main__':保护块里，否则会无限递归创建进程。这个坑在Linux上不存在，但从Linux迁移到Windows时经常遇到。

第四，pin_memory在CPU-only环境下没有意义，反而会增加内存开销。如果是在没有GPU的机器上做数据预处理测试，记得关掉。

7. 当常规手段不够用：DALI与自定义批处理

7.1 NVIDIA DALI的适用场景

如果CPU优化做到头了，GPU利用率还是上不去，可以考虑NVIDIA的DALI库。DALI把数据解码、增强、甚至部分预处理都放到GPU上执行，CPU只负责读取原始字节流。在图像和视频任务上，DALI通常能把数据加载吞吐提升3到5倍。

但DALI的引入成本不低：它有自己的算子体系，和torchvision.transforms不兼容，需要重写数据管道。而且DALI对数据格式有要求，不是所有数据集都能直接套用。我的建议是，只有在CPU优化已经做到极致、且数据加载确实是硬瓶颈的情况下才考虑DALI。

7.2 自定义batch sampler减少无效等待

在某些场景下，不同样本的预处理耗时差异很大（比如图片尺寸不一）。默认的RandomSampler随机抽取样本，可能导致某个batch里全是高耗时样本，拖慢整个batch。这时候可以用自定义的BatchSampler，把耗时相近的样本分到同一个batch里。

from torch.utils.data import Sampler class BucketBatchSampler(Sampler): def __init__(self, lengths, batch_size): # 按长度排序后分桶 self.buckets = [lengths[i:i+batch_size] for i in range(0, len(lengths), batch_size)] def __iter__(self): import random random.shuffle(self.buckets) for bucket in self.buckets: yield bucket def __len__(self): return len(self.buckets)

这种做法在NLP的变长序列任务里很常见，在图像任务里如果图片尺寸差异大也可以用。代价是打破了完全随机的采样，可能对收敛有轻微影响，需要根据实际情况权衡。

7.3 监控与持续调优

性能排查不是一次性的工作。数据集变了、模型改了、硬件换了，瓶颈可能就转移了。我习惯在训练脚本里加一个简单的吞吐量监控：

import time class ThroughputMonitor: def __init__(self, window=50): self.window = window self.times = [] def update(self, batch_size): self.times.append((time.time(), batch_size)) if len(self.times) > self.window: self.times.pop(0) def throughput(self): if len(self.times) < 2: return 0 total_samples = sum(t[1] for t in self.times) elapsed = self.times[-1][0] - self.times[0][0] return total_samples / elapsed if elapsed > 0 else 0

每个step调用update，每隔几十个step打印一次throughput。如果吞吐量突然下降，说明数据侧或者模型侧出现了变化，可以及时排查。

我在实际项目里踩过最深的坑，是一个图像分割任务里用了自定义的collate_fn做padding，结果collate_fn里有一个嵌套的Python循环，每个batch要跑几百毫秒。profiler显示collate_fn占了CPU时间的40%，但因为它不在__getitem__里，一开始完全没往那个方向想。后来把padding逻辑改成用torch.nn.utils.rnn.pad_sequence，耗时直接降到了几毫秒。这个经历告诉我，排查性能问题时不要预设“瓶颈一定在某个地方”，让profiler的数据说话，比凭经验猜测靠谱得多。