PyTorch 2.0.1 DataLoader性能逊于1.13，求提速解决方案

问题背景 #

PyTorch 2.0.1 对 DataLoader 的 worker 管理机制做了变更：默认采用线程池替代 1.13 版本的独立进程模式，导致在 num_workers 非 0 时数据加载速度显著下降，且必须使用 2.0.1 版本的场景下需要针对性优化。

可行提速方案 #

1. 强制切换回多进程 worker 模式 #

PyTorch 2.0+ 保留了多进程 worker 的配置选项，通过指定 multiprocessing_context 参数，让 DataLoader 恢复到 1.13 版本的进程管理逻辑，规避线程池的 GIL 瓶颈（尤其适用于 CPU 密集型的数据预处理场景）：

from torch.utils.data import DataLoader

dataloader = DataLoader(
    your_dataset,
    batch_size=your_batch_size,
    num_workers=32,
    pin_memory=True,
    multiprocessing_context='fork'  # Ubuntu 系统推荐使用 fork，Windows 可选 spawn
)

2. 优化线程池模式下的配置 #

若需保留线程池模式，可通过以下参数减少开销：

• 启用 persistent_workers=True：让 worker 线程在 epoch 之间保持存活，避免重复创建销毁的性能损耗
• 调整 prefetch_factor：增大预取数据量（默认值为 2），减少主线程等待时间
  示例代码：

dataloader = DataLoader(
    your_dataset,
    batch_size=your_batch_size,
    num_workers=16,  # 线程池模式下建议测试 8-16 的最优值
    pin_memory=True,
    persistent_workers=True,
    prefetch_factor=4
)

3. 优化数据集加载逻辑 #

• 替换 CPU 密集型预处理操作：优先使用 PyTorch 内置算子或 torchvision.transforms.v2 的优化版本，这类算子在多线程环境下能更好地规避 GIL 限制
• 减少磁盘 IO 开销：将常用数据集提前缓存到内存或 SSD，比如自定义 Dataset 类时在初始化阶段加载全部数据到内存，或使用内存映射文件读取大文件

4. 寻找最优 worker 数量 #

线程池模式下的最优 num_workers 与多进程模式不同，建议测试 8-16 的范围（而非原 32），找到当前机器的性能平衡点——过多线程会导致频繁切换，反而降低效率。

验证方式 #

保持其他参数与复现步骤一致，仅调整上述优化参数，统计数据加载的总耗时，对比性能提升效果。

内容的提问来源于stack exchange，提问作者Milad Sikaroudi