Pytorch指定GPU的方法總結

Pytorch指定GPU的方法

改變系統變量

1. 改變系統環境變量僅使目標顯卡,編輯 .bashrc文件，添加系統變量

   export CUDA_VISIBLE_DEVICES=0
   #這里是要使用的GPU編號
   

2. 在程序開頭設置

   os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3'
   

3. 在運行程序時指定

   # 運行程序時使用命令行,來設置該程序可見的gpu:
   CUDA_VISIBLE_DEVICES=0,1,2,3 python xxx.py
   

使用torch.cuda接口

# 在生成網絡對象之前執行
torch.cuda.set_device(0)

使用pytorch的並行GPU接口

model = torch.nn.DataParallel(model, device_ids=[0])

初始化模型時

model = Model.cuda(0)