如何修复Python程序中PyTorch多进程启动模式导致的死锁问题？

根本原因是子进程启动时复用父进程的CUDA上下文或OpenCV等非fork-safe库的全局状态导致初始化阻塞；典型表现为卡在for batch in dataloader:且无报错，CPU占用低、GPU显存无变化。

PyTorch DataLoader 设置 num_workers > 0 后卡在 for batch in dataloader:，CPU 占用低、GPU 显存不动、无报错——这不是代码写错了，而是子进程启动时继承了父进程的 CUDA 上下文或 OpenCV 等非 fork-safe 库的全局状态，导致初始化阻塞。

为什么 fork 启动方式在 Linux 下极易死锁

Linux 默认用 fork 启动子进程，它会复制父进程的内存镜像，包括已初始化的 CUDA 上下文、OpenCV 线程池、logging 配置等。这些状态在子进程中无法安全复用，尤其当主进程提前调用了 torch.cuda.is_available() 或 cv2.imread() 时，子进程可能卡在 futex 等待或 Bus error 中。

• fork 性能好但危险：只要父进程碰过 CUDA、OpenCV、matplotlib、h5py 或任何带后台线程的模块，就可能埋雷
• 死锁现象高度随机：可能第 1 个 epoch 就卡，也可能跑完 10 个 epoch 才崩，难以复现
• ps aux | grep python 可见若干 worker 进程状态为 D（不可中断睡眠）或长时间 R 但无日志输出

必须在 if __name__ == ‘__main__’: 中设 spawn

把 torch.multiprocessing.set_start_method('spawn') 放在脚本最顶部的 if __name__ == '__main__': 块内，且早于任何 DataLoader 创建和 CUDA 操作。

• 这是跨平台兼容解法：Windows/macOS 强制 spawn，统一设成它反而省心
• 必须加 force=True（尤其旧环境或混合 CUDA 场景）：torch.multiprocessing.set_start_method('spawn', force=True)
• 不能放在 Jupyter 或交互式环境中执行，会报 RuntimeError: context has already been set
• 设完之后，Dataset.__init__ 中不能传入不可序列化对象：比如 lambda、打开的文件句柄、threading.Lock、未 pickle 化的模型实例

配套关闭 pin_memory 和延迟 CUDA 初始化

pin_memory=True 会在子进程启动时触发 CUDA 上下文预分配，加剧 fork 冲突；而提前创建 device = torch.device('cuda') 或 torch.cuda.empty_cache() 也会污染子进程环境。

立即学习“Python免费学习笔记（深入）”；

• 调试阶段或 CPU-only 环境，显式设 pin_memory=False
• 所有 CUDA 相关操作（model.to('cuda')、tensor.cuda()、torch.cuda.is_available()）必须挪到 DataLoader 实例化之后
• 避免在 __getitem__ 中隐式触发 logging / matplotlib / h5py —— 它们不是 multiprocessing-safe 的
• 若用 OpenCV，可在 __getitem__ 开头加 cv2.setNumThreads(0)，或换装 opencv-python-headless

排查时别只看 Python 日志

死锁常发生在系统级同步原语上，Python 层面安静得可怕。真实线索藏在进程状态和系统调用里。

• 运行 ps aux | grep python，观察 worker 进程是否长期处于 D 或高 CPU R 状态
• 对卡住的 worker PID 执行 strace -p <pid></pid>，若看到大量 futex(..., FUTEX_WAIT, ...)，基本锁定是线程锁未重置
• 用 lsof -p <pid></pid> 查看子进程打开了哪些文件/共享库，辅助判断是否加载了冲突的 GUI 依赖（如 Qt、X11）
• CUDA_LAUNCH_BLOCKING=1 对这类死锁完全无效，它只捕获 kernel 同步错误，不干预 NCCL 或 fork 初始化问题

真正麻烦的不是改几行代码，而是那些“看似无关”的前置操作——比如某个 utils 模块顶层 import 时悄悄调用了 torch.cuda.device_count()，或者 logger 配置里启用了多线程 handler。这些细节一旦混进主模块导入链，就会在 spawn 模式下被每个子进程重复执行，变成隐形炸弹。