torch設置GPU

設置torch可見哪幾個GPU：

import os

os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"  

# 例子>> "0"：只有0號GPU可見；"1,3"：只有1號和3號GPU可見；""：全部不可見，此時 torch.cuda.is_available() 輸出 False。

注意：不一定非要寫在最前面，但至少必須要在 import torch 之前，否則可能不起作用。

 

使用torch.cuda.set_device()可以更方便地將模型和數據加載到對應GPU上, 直接定義模型之前加入一行代碼即可

torch.cuda.set_device(gpu_id) #單卡
torch.cuda.set_device('cuda:'+str(gpu_ids)) #可指定多卡

但是這種寫法的優先級低，如果model.cuda()中指定了參數，那么torch.cuda.set_device()會失效，而且pytorch的官方文檔中明確說明，不建議用戶使用該方法。