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

————————————————
版权声明:本文为CSDN博主「kongkongqixi」的原创文章,遵循CC 4.0 BY-SA版权协议,转载请附上原文出处链接及本声明。
原文链接:https://blog.csdn.net/kongkongqixi/article/details/100521590
————————————————

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

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

  3. 修改默认编号的方法,有两种方法:
    方法①:
    torch.cuda.set_device(1)
    进行指定默认gpu使用编号,同样也可以用
    torch.nn.DataParallel(model, device_ids=[1, 2])
    进行指定,也必须包含set_device(1)指定的device:1的设备,缺点是会占用一些device:0的gpu内存
    方法②:
    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])

你可能感兴趣的:(Python,pytorch,人工智能,python)