編程速記(7)：Pytorch篇-報錯：Invalid device ids - 多GPU並行

報錯

invalid device ids

場景

在運行多GPU的pytorch程序時，報該錯誤

原因

因爲pytorch 默認是隻有cuda 0 是可見的。可以運行下面代碼：

import torch
print(torch.cuda.device_count())

會發現服務器上明明有多張顯卡，但是返回值爲1.這是因爲：pytorch默認只有cuda0可見

解決方案

假定要執行的文件名爲：main_multi_gpu.py

CUDA_VISIBLE_DEVICE=0,1 python main_multi_gpu.py 

注意：上面的數字對應在程序中使用的gpu編號。