主要是针对指定了两片GPU,但默认使用第一块的问题。
我在用pytorch跑程序的时候,由于代码比较大,需要用两块板子。但无论是在代码中加入
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
还是 直接在命令行输入
CUDA_VISIBLE_DEVICES='0,1' python train.py
都无法解决问题。后来发现应使用 torch.nn.DataParallel这个函数,具体用法如下所示:
#自动选择GPU
model = nn.DataParallel(model)
model = model.cuda()
#选择指定GPU
model = torch.nn.DataParallel(model.cuda(),device_ids=[2,3])
#指定第二、三块板子
但是发现报错:AssertionError: Invalid device id,这是因为pytorch默认使用gpu编号为device:0的设备,对模型指定编号gpu训练时,必须要有编号为device:0的gpu,不然会报AssertionError: Invalid device id错误;所以如果你想使用2、3块板子,需要修改pytorch默认编号,故代码修改为:
torch.cuda.set_device(2)#指定默认编号为2
model = torch.nn.DataParallel(model.cuda(),device_ids=[2,3])
还有一种方法是:
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"
model = torch.nn.DataParallel(model.cuda(),device_ids=[2,3])
#pytorch 会默认把device:2改为device:0,而device:3变为device:1
具体可参考这篇博客.
问题解决!欧耶!