首先指定GPU,如下代码:
import os
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])
即0对应真实设备的1,1对应真实设备的2.注意,容易报错
下面开始配置使用多块gpu训练:
对于pyTorch的配置,要配置如下三处:模型,损失函数(或者优化器),数据。
device_ids=[0,1],(推荐建立一个列表存储)
1.模型
model = VisitNet().cuda(device_ids[0])
model = torch.nn.DataParallel(model, device_ids=device_ids)
2.损失函数
criterion = nn.CrossEntropyLoss().cuda(device_ids[0])
optimizer = nn.DataParallel(criterion, device_ids=device_ids)
3.数据
input = input.cuda(device_ids[0])
target = target.cuda(device_ids[0])