WSL下VSCode中PyTorch DataLoader设置num_workers>0时的CUDA初始化错误

问题现象 #

• 设置 num_workers=0 时，DataLoader运行正常：

  test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=0)
  

• 设置 num_workers≥1 时，抛出错误：

  RuntimeError: CUDA error: initialization error
  

数据集为MNIST，按50%/50%划分训练集与测试集。

环境配置 #

硬件 #

• GPU：NVIDIA GeForce RTX 3060 Ti（显存8191.50MB）
• CPU：物理核心12，逻辑核心24
• 内存：总15.44GB，可用12.71GB

软件 #

• 操作系统：Linux（WSL环境下的VS Code，OS版本#1 SMP Fri Mar 29 23:14:13 UTC 2024）
• Python版本：3.11.4
• PyTorch：2.0.1+cu118（检测到1块GPU）
• TensorFlow：2.14.0（检测到1块GPU）、Keras 2.14.0

解决建议 #

1. 配置WSL环境变量
   在WSL终端添加以下环境变量，避免子进程CUDA初始化冲突：

export CUDA_VISIBLE_DEVICES=0
export PYTHONFAULTHANDLER=1

可将上述内容添加到~/.bashrc或~/.zshrc中，确保每次启动WSL自动生效。

2. 调整DataLoader参数

• 启用pin_memory=True，帮助子进程高效将数据传输到GPU：

  test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=1, pin_memory=True)
  

• 控制num_workers数量：建议设置为CPU逻辑核心数的一半以内（比如≤12），减少资源竞争。

3. 规范CUDA上下文初始化时机
   避免在主线程创建DataLoader前提前初始化CUDA上下文（例如不要提前调用torch.cuda.device()或把张量移到GPU），让worker进程自行完成CUDA初始化。

4. 更新WSL与显卡驱动

• 将Windows端NVIDIA显卡驱动更新至最新版本，确保WSL2的GPU支持正常。
• 在Windows终端运行wsl --update，更新WSL内核，修复可能的GPU交互底层问题。

5. 隔离PyTorch与TensorFlow环境
   当前环境同时安装了PyTorch和TensorFlow，两者可能在CUDA上下文管理上存在冲突。建议创建独立虚拟环境分别运行两个框架的项目，避免资源抢占。

内容的提问来源于stack exchange，提问作者Kalin Stoyanov