编程速记(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编号。