PyTorch DataLoader `num_workers` 配置指南：从新手到进阶

​ 在 PyTorch 中，DataLoader 是训练流程中不可或缺的一环，它负责将数据集批量化并加载到模型中。在训练大型模型或处理图像、视频等大数据量任务时，数据加载速度往往会成为瓶颈。

​ num_workers 参数控制 DataLoader 使用的子进程数量，是优化数据加载性能的关键配置。本文将详细解析 num_workers 的取值意义、对训练性能的影响，以及在 Docker 环境下常见问题和解决方案。

1. num_workers 参数说明

​ 决定了 DataLoader 在加载数据时使用的进程数：

值	说明
0	数据在主进程中加载（单线程）
>0	使用指定数量的子进程并行加载数据

⚠️ 注意：num_workers=0 时，数据加载简单可靠，但可能成为 GPU 的瓶颈；而 num_workers>0 可以加速数据加载，但可能占用大量内存。

2. 不同取值的影响

(1) num_workers=0

• 优点：无需额外内存，兼容性好，调试方便

• 缺点：数据加载速度慢，GPU 可能需要等待数据，对大数据集训练效率低

train_loader = DataLoader(dataset, batch_size=64, num_workers=0)

(2) num_workers>0（如 4、8）

• 优点：并行加载数据，提高训练速度，减少 GPU 等待时间

• 缺点：占用更多内存和共享内存，可能导致内存问题

  train_loader = DataLoader(dataset, batch_size=64, num_workers=4, pin_memory=True)

💡 Tip：结合 pin_memory=True 可以加速 GPU 数据传输，但需要系统内存充足。

3. 常见问题：共享内存不足错误

​ 在 Docker 中使用 num_workers>0 时，可能出现如下错误：

RuntimeError: DataLoader worker (pid xxx) is killed by signal: Bus error.
It is possible that dataloader's workers are out of shared memory.
Please try to raise your shared memory limit.

​ 原因：Docker 容器默认共享内存（/dev/shm）只有 64MB，多个 worker 进程需要使用共享内存进行数据传输，内存不足时会导致 Bus error。

4. 解决方案

方案一：增加 Docker 共享内存（推荐）

​ 在 docker-compose.yml 中添加 shm_size 配置：

services:
  your_service:
    image: your_image
    shm_size: 2g    # 设置共享内存为 2GB
    # ... 其他配置

​ 或使用 docker run 命令时添加参数：

docker run --shm-size=2g your_image

• 运行后可以在容器中执行：

  df -h /dev/shm      

  验证 shm_size 配置是否生效。

  Filesystem      Size  Used Avail Use% Mounted on
  shm             2.0G     0  2.0G   0% /dev/shm

方案二：减少 num_workers

​ 当内存有限时，将 num_workers 调小或设为 0：

# 完全禁用多进程加载
DataLoader(dataset, batch_size=64, num_workers=0)

# 或使用较少的 worker
DataLoader(dataset, batch_size=64, num_workers=2)

方案三：使用 pin_memory 优化

在 CPU 内存充足时，结合 pin_memory=True 可以提高 GPU 数据传输速度：

DataLoader(dataset, batch_size=64, num_workers=4, pin_memory=True)

5. 推荐配置指南

环境	num_workers	备注
Docker（默认 shm）	0	避免共享内存不足问题
Docker + shm_size=2g	4-8	可根据 CPU 核心数调整
本地机器	CPU 核心数	一般设置 4-8
调试阶段	0	方便排查问题

​ 小技巧：在训练前用小批量数据跑一个 DataLoader 测试，观察 CPU、GPU 和内存占用，找到最适合的 num_workers 配置。

6. 总结

​ num_workers 对训练性能影响巨大：

• 0：安全、易调试，但速度慢
• >0：可加速数据加载，但需注意内存与 Docker 共享内存限制

​ 在 Docker 中使用多进程 DataLoader 时，优先考虑增加 shm_size ，或降低 num_workers，结合 pin_memory 可以进一步优化 GPU 数据传输效率。

7. 参考资料

• PyTorch DataLoader 文档
• Docker 共享内存配置