在Python中怎么解决PyTorch多进程DataLoader死锁问题-Python教程-PHP中文网

dataloader在num_workers>0时死锁的根本原因是子进程fork时继承了主进程未重置的cuda上下文，导致驱动层资源争用；常见于dataset.__init__、全局作用域或__getitem__中误触发gpu操作。

在python中怎么解决pytorch多进程dataloader死锁问题

为什么 `num_workers > 0` 时 DataLoader 容易死锁

根本原因不是 PyTorch 本身有 bug，而是子进程在初始化时意外触发了主进程才该做的操作——比如调用 torch.cuda.set_device()、访问已关闭的文件句柄、或导入含全局 CUDA 初始化逻辑的模块。Linux 下 fork 模式会复制主进程内存状态，若此时 CUDA 上下文未被正确重置，子进程可能卡在驱动层等待资源，表现为“没报错但训练不动”。常见于 Jupyter 环境、多卡训练脚本开头误写 torch.cuda.set_device()，或自定义 Dataset 的 __init__ 中做了 GPU 相关操作。

必须检查的三个初始化位置

死锁往往藏在你看不见的地方。重点排查：

Dataset.__init__：禁止在此处创建 torch.Tensor、调用 cuda()、或初始化任何依赖 CUDA 的对象（如某些预处理库的 GPU 加速器）
全局作用域：脚本最外层（非函数内）写了 torch.cuda.set_device(0) 或 model.to('cuda') —— 这会在每个 worker fork 后重复执行，引发竞争
__getitem__ 内部：看似安全，但如果用了第三方库（如 torchaudio 0.12+ 默认启用 CUDA 后端），可能隐式触发设备初始化

快速验证和绕过方案

先确认是不是 num_workers 导致的：把 DataLoader 的 num_workers=0 运行一次，如果正常，基本锁定问题范围。然后按顺序尝试：

Secure Python Code Sandbox

安全 Python 代码沙箱：在隔离的沙箱中执行任意 Python 代码，预装 requests、NumPy、pandas 库。返回 stdout、stderr。

下载

加 pin_memory=False：避免 pin memory 时触发 CUDA 上下文创建（尤其 Windows + WSL 环境）
换启动方式：把 spawn 替代默认 fork，在 DataLoader 外加 torch.multiprocessing.set_start_method('spawn', force=True)（注意只能调用一次，且需放在 if __name__ == '__main__': 下）

限制 worker 初始化：用 worker_init_fn 显式控制每个 worker 的行为，例如：

def worker_init_fn(worker_id):
    # 清除可能残留的 CUDA 状态
    import torch
    torch.cuda.empty_cache()
    # 避免所有 worker 绑定同一 GPU
    torch.cuda.set_device(worker_id % torch.cuda.device_count())

再传入 DataLoader(worker_init_fn=worker_init_fn)

Windows 和 macOS 用户的特殊陷阱

这两个系统不支持 fork，PyTorch 默认强制用 spawn，但 spawn 模式下主进程的全局变量、CUDA 状态不会继承，反而容易因模块重载失败导致静默卡死。典型表现是：训练卡在第一个 batch，nvidia-smi 看不到新进程显存占用，ps aux | grep python 却能看到多个 worker 进程处于 S（sleep）状态。此时必须确保：

所有数据加载逻辑封装在函数内，不要依赖全局模型/设备变量
自定义 Dataset 类不能在类定义体中执行 GPU 操作（比如 self.transform = SomeGPULib()）
升级到 PyTorch ≥ 1.12，并在 DataLoader 中显式指定 persistent_workers=True 减少反复 spawn 开销

真正棘手的是那些不报错、不超时、只在特定 batch 数后卡住的情况——这往往意味着某个 worker 在读取某个特殊样本时触发了未捕获异常，而异常被 multiprocessing 吞掉了。建议在 __getitem__ 最外层加 try...except 并打印 worker_id 和 index，定位具体哪条数据出问题。