运行时错误：CUDA 错误 - PyTorch DataLoader 工作进程中的初始化错误

Question

我在使用 DataLoader 进行训练时在 PyTorch 代码中遇到运行时错误。错误发生在工作进程中，回溯指向 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.

# Code snippet where the error occurs
for batch, (X, Y) in enumerate(train_dataloader): <---- error here 
    print(f"Batch: {batch+1}")
    print(f"X shape: {X.shape}")

我正在使用 PyTorch 进行深度学习项目，并且仅在我尝试迭代数据加载器时才会发生。该错误似乎与 CUDA 初始化有关，但我不知道如何解决它。

我正在 Google Colab 上运行：
PyTorch 版本：2.1.0+cu121
GPU型号：Tesla T4

我使用以下代码设置设备：

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

它确实会抢占 GPU。我在初始化时将设备交给张量，这发生在测试/训练分割函数中。

我尝试使用

os.environ["TORCH_USE_CUDA_DSA"] = "1"

设置 TORCH_USE_CUDA_DSA，但没有成功。 该错误是什么意思/为什么会发生？当我的手电筒设备设置为“cpu”时，它工作正常，因此整体数据摄取/训练不是问题。

数据类非常简单：

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

我在设置数据加载器时考虑了 GPU 的使用：

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
)


# Check it's working
for batch, (X, Y) in enumerate(train_dataloader): <--- error
    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

奇怪的是，当我删除 num_workers 和 pin_memory 参数时，它起作用了。这些对于提高训练运行时间很有价值，但也许我误用了它们？不确定那里有什么交易。

有办法设置TORCH_USE_CUDA_DSA吗？

Answer 1

要查看错误来自何处，请通过设置 device = torch.device("cpu") 在 cpu 上运行它。解决问题后您可以将设备切换到cuda