torch单机多卡重点：

 

针对于单机多卡，可以使用 nn.DataParallel 进行训练，需要注意的是，与单卡训练有一些地方不同：

（1）模型和优化器的初始化：

模型初始化之后，需要传入 nn.DataParallel ，再进行并行化处理，同时注意优化器同样需要做并行化处理，如下：

model = Model.Model(param).cuda()

model = nn.DataParallel(model, device_ids=device_ids) # multi-GPU 

optimizer = torch.optim.Adam(model.parameters(), lr=0.002)  # 优化器

optimizer = nn.DataParallel(optimizer, device_ids=device_ids) # multi-GPU 

 

device_ids是一个list，保存要使用的GPU标号：

torch.cuda.set_device(1)

device_ids = [1, 2]

 

（2）优化器参数的更新：

optimizer.module.step()

与常规相比，多了个.module，是因为在并行化处理的时候，torch也对模型参数的结构进行了外部的封装，因此需要多一层才能获取到参数

 

（3）模型的保存：

torch建议只保存参数，可移植性更强：

torch.save(model.module.state_dict(), root_path+'data/model/train: epoch:'+str(epoch+1)+',step:'+str(step)+',loss:'+str(min_train_loss)[:6]+'.pth')

还有就是模型整个保存，官方并不推荐，先不展示了

 

（4）模型的加载：

这里只展示保存的参数的加载，也就是将上面方法保存的模型进行加载：

# 加载模型Model
model = Model.Model(param)
model = nn.DataParallel(model, device_ids=[1, 2]).cuda() # multi-GPU 
# 获得模型参数
model_dict = torch.load('/home/xxx/project/xxx/data/model/train: epoch:1,step:266,loss:1.4726.pth')
# 载入参数
model.module.load_state_dict(model_dict)

 

其他：

 nn.DataParallel  可以实现单机多卡的训练，但是模型的保存也分为很多种，比如单机多卡训练出的模型，如何在无GPU的环境下加载，如何在单卡情况下加载等等…

然后，还有一些分布式训练（多机多卡）的方法也可以用在单机多卡的情况下，这些还没有进行尝试。