【實用】多GPU未占滿情況下，使用Python指定GPU運作代碼

我們有時候會遇到一種情況，即伺服器上某節點包含多個GPU，我們已經使用了一個或者幾個GPU，還有部分GPU未使用，這個時候我們再次将代碼上傳伺服器進行運作，發現不知道如何進行選擇特定的GPU運作，且如果隻通過預設運作，會發現報錯，因為程式預設是從 GPU 0 開始的，如果GPU 0被占用了，則無法運作其他指令，會報 out of CUDA memory的異常。

遇到以上的問題，其實有兩個解決方案，如下所示：

1 使用 nvidia-smi 進行檢視GPU的資訊，發現好多空載的GPU，這裡注意每個GPU前面的編号 (0,1,2,…5)，這個就是後續指定GPU編号使用的。

1 在Python代碼中直接聲明

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "3,4"#即指定使用 3 4 号GPU
           

import torch
#這個id 貌似隻能是單個整數
torch.cuda.set_device(id)
           

2 在運作指令中標明GPU

在指令中使用 CUDA_VISIBLE_DEVICES=1,2 進行標明，如下，例子中為指定1，2兩塊GPU進行運作該程式

CUDA_VISIBLE_DEVICES=1,2 python -m torch.distributed.launch --nproc_per_node 2 --master_port 12345  main.py 
           

以上就是我的解決過程，希望可以幫助到您