pytorch 模型训练的多GPU训练

pytorch针对多GPU训练提供了两个类来实现多GPU训练，分别是torch.nn.DataParallel和torch.nn.parallel.DistributedDataParallel，因为torch.nn.DataParallel相对比较简单，因此先介绍torch.nn.DataParallel。

当使用torch.nn.DataParallel的时候，数据会分配到各个GPU中，完成前向计算，然后再汇总到主卡计算loss，进行反向传播。显然，这种方式会使得多个GPU之间负载不均衡。

torch.nn.DataParallel具体的使用方法

model = BERTModel(pretrain_name,vocab_size)
device_ids = [0,1] #指定使用的显卡id
model = nn.DataParallel(model,device_ids=device_ids).to(device_ids[0]) # 多卡训练，用device_ids[0]只是定义一个样式，不需要逐卡指定设备。但没这一步会报错。数据的batch_size为单卡的batch_size*len(devices)

trainer = torch.optim.Adam(model.parameters(),lr=1e-3)
trainer = nn.DataParallel(trainer,device_ids=device_ids)

注意device_ids 为指定需要训练的GPU，以上的几行代码可以实现模型和优化器都放到GPU中，记得在进行训练的时候，数据也要放到GPU中，只需要将数据的tensor.cuda(device_ids[0]) 即可。torch.nn.DataParallel模块会自动将数据都分发到各个GPU中进行计算。