Pytorch 多 GPU 問題

Question

我想用 2 個 GPU(id 5, 6) 訓練我的模型，所以我用CUDA_VISIBLE_DEVICES=5,6 train.py運行我的代碼。 但是，當我打印 torch.cuda.current_device 時，我仍然得到 id 0而不是 5,6。 但是 torch.cuda.device_count 是2 ，這似乎是正確的。 如何正確使用 GPU5,6？

Answer 1

這很可能是正確的。 PyTorch 只能看到兩個 GPU（因此索引為 0 和 1），它們實際上是您的 GPU 5 和 6。

使用nvidia-smi檢查實際使用情況。 如果仍然不一致，您可能需要設置一個環境變量：

export CUDA_DEVICE_ORDER=PCI_BUS_ID

（請參閱'nvidia-smi -L' 和 cuDeviceGetName() 之間的 ID 不一致）

Answer 2

您可以檢查設備名稱以驗證它是否是該 GPU 的正確名稱。 但是，我認為當您將 Cuda_Visible 設置在外面時，您已經迫使手電筒只看那 2 個 gpu。 所以torch會手動為它們設置index為0和1。因此，當你檢查current_device時，它會輸出0