PyTorch DataLoader多进程CUDA初始化错误排查求助

错误现象 #

在Google Colab（PyTorch 2.1.0+cu121，Tesla T4 GPU）中训练时，遍历DataLoader触发CUDA初始化错误，报错信息如下：

RuntimeError: Caught RuntimeError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/_utils/worker.py", line 308, in _worker_loop
    data = fetcher.fetch(index)
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/_utils/fetch.py", line 51, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/_utils/fetch.py", line 51, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "<ipython-input-38-b0835867e243>", line 8, in __getitem__
    return self.X[index], self.Y[index]
RuntimeError: CUDA error: initialization error
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

出错代码片段：

# 错误触发位置
for batch, (X, Y) in enumerate(train_dataloader): 
    print(f"Batch: {batch+1}")
    print(f"X shape: {X.shape}")

自定义Dataset类：

class Data(Dataset):
    def __init__(self, X, Y):
        self.X = X
        self.Y = Y
        self.len = len(X)

    def __getitem__(self, index):
        return self.X[index], self.Y[index]

    def __len__(self):
        return self.len

DataLoader初始化代码：

batch_size = 256

train_data = Data(X_train, Y_train)
train_dataloader = DataLoader(
  dataset=train_data, 
  batch_size=batch_size, 
  shuffle=True,
  pin_memory=True
)

# 测试代码触发错误
for batch, (X, Y) in enumerate(train_dataloader):
    print(f"Batch: {batch+1}")
    print(f"X shape: {X.shape}")
    print(f"y shape: {Y.shape}")
    print(f"X type: {X.dtype}")
    print(f"Y type: {Y.dtype}")
    break

已知条件：已指定GPU设备，初始化阶段数据已移至GPU；移除num_workers和pin_memory后程序正常，但需要这两个参数提升效率；尝试设置os.environ["TORCH_USE_CUDA_DSA"] = "1"无效。

问题原因 #

1. 子进程无法访问主进程CUDA张量：DataLoader的worker进程是独立子进程，无法直接读取主进程中已移至GPU的张量，读取时会触发CUDA初始化错误。
2. pin_memory参数误用：pin_memory=True的作用是锁定CPU内存页，加速CPU到GPU的数据传输，但如果数据已经在GPU上，该参数不仅无效，还会引发内存访问冲突。

解决方法 #

方法1：调整数据存储位置，在合适时机转移到GPU #

方案A：在Dataset的__getitem__中转移数据

让数据默认存在CPU，获取单条数据时再转到GPU：

class Data(Dataset):
    def __init__(self, X, Y):
        # 确保数据在CPU存储，若之前已移到GPU则转回
        self.X = X.cpu() if X.is_cuda else X
        self.Y = Y.cpu() if Y.is_cuda else Y
        self.len = len(X)

    def __getitem__(self, index):
        return self.X[index].to(device), self.Y[index].to(device)

    def __len__(self):
        return self.len

方案B：在训练循环中统一转移数据

保持Dataset不变，在训练时批量转移数据到GPU：

for batch, (X, Y) in enumerate(train_dataloader):
    X = X.to(device)
    Y = Y.to(device)
    # 后续训练逻辑

方法2：正确配置DataLoader参数 #

保留num_workers（根据Colab资源设置为2-4即可），pin_memory保持True（此时数据在CPU，锁页能加速转GPU）：

train_dataloader = DataLoader(
  dataset=train_data, 
  batch_size=batch_size, 
  shuffle=True,
  num_workers=2,
  pin_memory=True
)

TORCH_USE_CUDA_DSA正确设置方式 #

该环境变量必须在导入PyTorch前设置才会生效，在Colab中要放在代码最开头：

import os
os.environ["TORCH_USE_CUDA_DSA"] = "1"
# 之后再导入PyTorch相关模块
import torch
from torch.utils.data import Dataset, DataLoader

注：开启DSA仅用于获取更详细的设备端错误信息，解决当前问题的核心还是调整数据存储位置。

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