PyTorch DataLoader在num_workers>0时卡住,主因是spawn启动方式下子进程重复导入主模块并执行GPU初始化等副作用操作;必须用if name == '__main__':包裹数据加载逻辑,且Dataset中禁止任何GPU相关操作。

PyTorch 的 DataLoader 在 num_workers > 0 时卡住不动,多数情况不是代码写错,而是子进程初始化阶段就阻塞了——尤其常见于 Windows、Jupyter 或 import 冲突场景。
为什么 DataLoader 多进程会死锁?
根本原因在于 Python 多进程的 spawn 启动方式(Windows/macOS 默认)会重新导入主模块,若主模块顶层执行了 GPU 初始化、模型构建或依赖全局状态的操作,子进程可能重复调用 torch.cuda.init() 或等待未释放的资源锁。Jupyter 环境下更危险:内核已加载过 CUDA 上下文,子进程 fork/spawn 时状态不一致。
-
num_workers=0正常但设为1就卡住 → 基本锁定是子进程启动失败 - 卡在
next(iter(dataloader))第一次调用 → 阻塞发生在 worker 进程首次初始化 - Linux 下偶尔正常、Windows/Jupyter 必现 → 启动方式差异暴露了隐式依赖
必须加 if __name__ == '__main__': 包裹入口
这是 spawn 模式下的硬性要求。没它,子进程导入脚本时会再次执行训练循环、模型定义甚至 torch.device('cuda'),触发 CUDA 上下文冲突。
if __name__ == '__main__':
dataset = MyDataset()
dataloader = DataLoader(dataset, num_workers=4, batch_size=32)
for batch in dataloader:
# 训练逻辑
pass- 哪怕你只写了个测试脚本,也得套这个 guard;Jupyter 用户需改用
if 'get_ipython' in globals():+ 显式启动逻辑 - 把
DataLoader创建和迭代放在函数里没用——关键是要让子进程启动时不执行任何副作用代码 - 第三方库(如 timm)内部调用
torch.hub.load也可能触发隐式 import,需确认其是否在顶层执行
避免在 Dataset.__init__ 或 __getitem__ 中初始化 GPU 资源
子进程无法共享父进程的 CUDA 上下文,任何在数据加载路径中调用 torch.cuda.device、torch.randn(..., device='cuda') 或加载预训练模型的操作都会导致死锁或 RuntimeError。
立即学习“Python免费学习笔记(深入)”;
- 把模型、tokenizer、图像变换器等全部移到
__init__外部,在__getitem__中只做纯 CPU 操作(PIL 打开、numpy 转换、torch.tensor 构造) - 如果必须用 GPU 加速预处理(极少见),改用
torch.compile或提前缓存结果,不要在 worker 进程里动态创建device='cuda'张量 - 检查自定义
collate_fn:它运行在主进程,但若引用了 Dataset 实例中的 GPU 属性,也会间接触发问题
调试死锁的最小化手段
别靠猜,用最简复现 + 系统级观察:
- 临时把
num_workers设为1,再加pin_memory=False和prefetch_factor=2排除显存映射干扰 - Linux 下运行
strace -f -e trace=clone,wait4 python train.py 2>&1 | grep -E '(clone|wait)'看子进程是否卡在 fork 后 - Windows 下用 Process Explorer 查看子进程状态:若显示 “Suspended” 或无 CPU 占用,基本是卡在 import 阶段
- 把
DataLoader的worker_init_fn设为lambda x: print(f'Worker {x} started'),看是否根本没输出
真正棘手的死锁往往藏在第三方包的 __init__.py 里——比如某个库在导入时自动调用 torch.backends.cudnn.benchmark = True,这种必须查源码或换版本。


















