Python多进程训练中PyTorch如何正确设置DataLoader的num

num_workers > 0 会卡住或报 pickleerror 是因为子进程需 pickle 反序列化 dataset、collate_fn 等，若含不可序列化对象（如 videocapture、lambda、动态 import、全局变量）或 windows 下未加 if name == '__main__': 保护导致重复执行，则失败。

python多进程训练中pytorch如何正确设置dataloader的num_workers？

num_workers 设为 0 是最安全的起点，尤其在 Windows 或使用 fork 多进程启动方式时；设为大于 0 时必须确保数据集、collate_fn 和 transform 都是可序列化的，且不能依赖全局状态或未初始化的模块。

为什么 num_workers > 0 会卡住或报 PickleError？

PyTorch DataLoader 在子进程中通过 pickle 重建 dataset、sampler、collate_fn 和 transform。一旦其中任意一个对象包含不可序列化的内容（比如 OpenCV VideoCapture 实例、lambda 函数、带闭包的局部函数、未 import 的模块引用），子进程就无法反序列化，导致卡死或抛出 PickleError。

常见踩坑点：

dataset.__init__() 中直接创建了 cv2.VideoCapture 或 h5py.File 对象
用 lambda x: x + 1 当作 transform，而 lambda 无法被 pickle
在 __getitem__ 中动态 import 某个模块（import 在子进程中可能失败或路径不同）
使用了依赖全局变量的函数（如 global_cfg.batch_size），但该变量在子进程中未定义

Windows 下 num_workers > 0 必须加 if name == 'main':

Windows 默认使用 spawn 启动子进程，不共享父进程的运行上下文。若主模块未加保护，子进程会重复执行整个脚本，导致反复初始化模型、重载数据、甚至递归启动新进程——最终卡死或 OOM。

正确写法示例：

Forge — Python Developer

Python开发助手 — 自动化脚本、API封装、命令行工具、数据管道、网页爬虫。经过测试的代码，包含错误处理和文档。当你需要...

下载

if __name__ == '__main__':
    train_loader = DataLoader(dataset, batch_size=32, num_workers=4)
    for epoch in range(10):
        for batch in train_loader:
            # 训练逻辑
            pass

注意：if __name__ == '__main__': 必须包裹所有涉及 DataLoader 初始化和迭代的代码，不只是 DataLoader 创建那一行。

num_workers 设置多少才合适？

没有固定最优值，需结合 CPU 核心数、I/O 延迟、batch_size 和数据预处理开销实测。经验性建议：

从 num_workers=0 跑通逻辑，再逐步增加
Linux/macOS 上，num_workers 通常设为 min(8, os.cpu_count() - 1)，避免抢占训练主进程资源
SSD + 简单 transform（如 ToTensor）时，num_workers=2~4 往往已足够；若用了 heavy augment（如 Albumentations 中的 large-scale warp），可尝试 num_workers=6~8
观察 train_loader 迭代时 GPU 利用率是否持续 >90%：如果经常掉到 0%，说明数据加载成为瓶颈，可增 num_workers；如果增加后 GPU 利用率没提升，反而 CPU 占用飙高或内存暴涨，说明已过载