采用nn.DataParallel实现多gpu训练
1、首先选择GPU训练,定义device,“cuda:0"代表起始的device_id为0,如果直接使用"cuda”,同样默认起始为0。
device = torch.device("cuda:0" if torch.cuda.is_available() else)
2、然后选择GPU的数量
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model, device_ids=[0,1,2])
model.to(device)