4、数据加载压测：DataLoader 性能分析、num_workers 调优、prefetch_factor 设置、内存映射文件（mmap）的使用

数据加载，说白了就是训练管道的「进水口」。

很多同学把精力都花在模型上，结果一跑起来，GPU 使用率才 30%。

我见过太多项目，模型写得漂亮，但数据加载成了瓶颈。GPU 空转，电费照付，你说亏不亏？

这一章，咱们就把 DataLoader 的底裤扒干净。

核心观点：数据加载的速度，必须快过 GPU 的计算速度。否则，GPU 就是在等你。

4.1 DataLoader 性能分析：瓶颈到底在哪？

先问个问题：你怎么知道数据加载是瓶颈？

很简单，看 GPU 利用率。如果 GPU 利用率低于 90%，而 CPU 跑满了，那八成是数据加载的问题。

我习惯用 nvidia-smi 和 htop 两个命令同时看。一个看 GPU，一个看 CPU。

我的小技巧：写个简单的计时器，在训练循环里记录 time.time()，看看每次 next(iter(dataloader)) 花了多久。如果超过 10ms，就得优化了。

DataLoader 的瓶颈通常有三个地方：

磁盘 I/O 慢：尤其是机械硬盘，或者网络文件系统。
预处理太重：比如图像解码、数据增强，这些都在 CPU 上跑。
进程间通信开销：worker 进程把数据传回主进程，也是有成本的。

我记得有一次，帮一个团队排查问题。他们的数据是 4K 高清图片，每张 10MB 以上。结果 DataLoader 的 worker 全在等磁盘，GPU 利用率不到 20%。后来换成 SSD，直接飙到 85%。

4.2 num_workers 调优：不是越多越好

num_workers 这个参数，很多人以为越大越好。其实不是。

它的作用是启动多个子进程，并行加载数据。但每个进程都要占用 CPU 和内存。

我建议的调优方法：

从 0 开始：num_workers=0 表示主进程自己加载，适合调试。
逐步增加：每次加 2，观察 GPU 利用率的变化。
上限是 CPU 核心数：一般设为 os.cpu_count() 或略少。

import torch
import os

# 我常用的配置
num_workers = min(8, os.cpu_count())  # 不超过 8 个
dataloader = torch.utils.data.DataLoader(
    dataset,
    batch_size=64,
    shuffle=True,
    num_workers=num_workers,
    pin_memory=True
)

注意：num_workers 太多会导致内存爆炸。每个 worker 都会复制一份数据。我曾经见过有人设了 32 个 worker，结果 64GB 内存直接爆了。

还有一个坑：在 Windows 上，num_workers 大于 0 可能会报错。这是因为 Windows 的进程管理机制不同。解决方案是加一句 if __name__ == '__main__':。

4.3 prefetch_factor 设置：让数据「提前就位」

prefetch_factor 是 PyTorch 1.7 之后引入的参数。它控制每个 worker 预取多少个 batch。

默认值是 2。什么意思呢？就是每个 worker 会提前加载 2 个 batch 的数据，放在队列里等着。

为什么要这么做？

你想想看，GPU 算完一个 batch，马上就能拿到下一个。不用等 worker 现去加载。

我一般这样设置：

数据加载快（比如 SSD + 小图片）：prefetch_factor=2 就够了。
数据加载慢（比如网络文件系统）：可以试试 prefetch_factor=4 或 6。
不要超过 8：否则内存占用会很高。

dataloader = torch.utils.data.DataLoader(
    dataset,
    batch_size=64,
    num_workers=4,
    prefetch_factor=4,  # 每个 worker 预取 4 个 batch
    pin_memory=True
)

我的经验：prefetch_factor 和 num_workers 要配合调。比如 num_workers=4，prefetch_factor=4，那总共预取了 16 个 batch。如果你的 batch 很大（比如 256 张图），内存压力会很大。

4.4 内存映射文件（mmap）：磁盘 I/O 的终极解法

mmap 是什么？说白了，就是把磁盘文件映射到内存地址空间。程序读写文件，就像读写内存一样。

好处是什么？

省去系统调用：不用 read/write 了，直接内存访问。
按需加载：只加载你访问的那部分数据，不是整个文件。
共享内存：多个进程可以共享同一块映射，省内存。

在 PyTorch 里，怎么用 mmap？

如果你用 torch.utils.data.Dataset，可以在 __init__ 里把数据映射进来。

import mmap
import os

class MmapDataset(torch.utils.data.Dataset):
    def __init__(self, data_path):
        self.file = open(data_path, 'rb')
        # 映射整个文件
        self.mmap = mmap.mmap(
            self.file.fileno(),
            0,  # 0 表示映射整个文件
            access=mmap.ACCESS_READ
        )
        # 假设每个样本是固定大小
        self.sample_size = 1024 * 1024  # 1MB
        self.num_samples = len(self.mmap) // self.sample_size
    
    def __len__(self):
        return self.num_samples
    
    def __getitem__(self, idx):
        start = idx * self.sample_size
        end = start + self.sample_size
        data = self.mmap[start:end]
        # 这里可以转成 tensor
        return torch.frombuffer(data, dtype=torch.uint8)

注意：mmap 适合大文件、随机访问的场景。如果文件很小，或者你总是顺序读取，mmap 的优势不明显。

我曾经在一个 NLP 项目里用过 mmap。数据是 200GB 的文本文件，每行一个样本。用 mmap 之后，加载时间从 5 分钟降到了 10 秒。因为操作系统只加载了实际访问的那几页。

4.5 综合压测：一个完整的测试脚本

说了这么多，不如动手跑一跑。下面是我常用的压测脚本：

import torch
import time
import os

def benchmark_dataloader(dataloader, num_batches=100):
    start = time.time()
    for i, batch in enumerate(dataloader):
        if i >= num_batches:
            break
        # 模拟 GPU 计算
        _ = batch[0].cuda()
    end = time.time()
    return (end - start) / num_batches

# 测试不同配置
configs = [
    {'num_workers': 0, 'prefetch_factor': 2},
    {'num_workers': 2, 'prefetch_factor': 2},
    {'num_workers': 4, 'prefetch_factor': 2},
    {'num_workers': 4, 'prefetch_factor': 4},
    {'num_workers': 8, 'prefetch_factor': 2},
]

for config in configs:
    loader = torch.utils.data.DataLoader(
        dataset,
        batch_size=64,
        shuffle=True,
        **config,
        pin_memory=True
    )
    avg_time = benchmark_dataloader(loader)
    print(f"Config {config}: avg time per batch = {avg_time:.3f}s")

跑完之后，你会看到不同配置的差异。我建议你记录下结果，做成表格：

num_workers	prefetch_factor	平均每 batch 耗时	GPU 利用率
0	2	0.125s	45%
2	2	0.068s	72%
4	2	0.042s	88%
4	4	0.038s	91%
8	2	0.035s	93%

看到没？从 0 到 4 个 worker，性能提升很明显。但再往上加，收益就递减了。

总结一下：

num_workers 从 0 开始，逐步增加，直到 GPU 利用率饱和。
prefetch_factor 默认 2，数据加载慢时加到 4。
mmap 适合大文件随机访问，能大幅降低 I/O 延迟。
别忘了用 pin_memory=True，能加速 CPU 到 GPU 的数据传输。

嗯，数据加载这块，其实还有很多细节。比如 persistent_workers、collate_fn 的优化等等。但上面这几个，已经能解决 80% 的问题了。

下次你跑训练，记得先看看 GPU 利用率。如果低于 90%，回来翻翻这一章。

📚 课程目录

4、数据加载压测：DataLoader 性能分析、num_workers 调优、prefetch_factor 设置、内存映射文件（mmap）的使用

4.1 DataLoader 性能分析：瓶颈到底在哪？

4.2 num_workers 调优：不是越多越好

4.3 prefetch_factor 设置：让数据「提前就位」

4.4 内存映射文件（mmap）：磁盘 I/O 的终极解法

4.5 综合压测：一个完整的测试脚本