pytorch关于多块gpu使用总结，报错AssertionError: Invalid device id

转载自：https://blog.csdn.net/kongkongqixi/article/details/100521590

pytorch默认使用gpu编号为device:0的设备，可以使用
torch.nn.DataParallel(model, device_ids=[0, 1])
对模型进行制定gpu指定编号多gpu训练，必须要有编号为device:0的gpu，不然会报AssertionError: Invalid device id错误；

当gpu编号为device:0的设备被占用时，指定其他编号gpu使用torch.nn.DataParallel(model, device_ids=[1, 2])指定gpu编号会出现AssertionError: Invalid device id错误，这是因为pytorch默认使用gpu编号为device:0的设备，需要修改pytorch默认编号；
修改默认编号的方法，有两种方法：

（1）使用
torch.cuda.set_device(1)
进行指定默认gpu使用编号，这样可以使用
torch.nn.DataParallel(model, device_ids=[1, 2])
进行指定，也必须包含set_device(1)指定的device:1的设备，缺点是仍然会存在占用一些device:0的gpu内存；

（2）使用

os.environ["CUDA_VISIBLE_DEVICES"] = "1,2"
进行指定使用设备，这样会修改pytorch感受的设备编号如上则把device:1改为device:0，device:2改为device:1，则pytorch感知的编号还是从device:0开始，指定使用程序为：

torch.nn.DataParallel(model, device_ids=[0,1])
————————————————
版权声明：本文为CSDN博主「kongkongqixi」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/kongkongqixi/article/details/100521590