最近训练模型时候想要使用使用多GPU运算来提高计算速度，参考一些博客以及自己的动手实验搞懂了Pytorch的Multi-GPU原理。现在稍微整理一下。
原理
通常情况下，多GPU运算分为单机多卡和多机多卡，两者在pytorch上面的实现并不相同，因为多机时，需要多个机器之间的通信协议等设置。
pytorch实现单机多卡十分容易，其基本原理就是：加入我们一次性读入一个batch的数据, 其大小为[16, 10, 5]，我们有四张卡可以使用。那么计算过程遵循以下步骤：
- 假设我们有4个GPU可以用，pytorch先把模型同步放到4个GPU中。
- 那么首先将数据分为4份，按照次序放置到四个GPU的模型中，每一份大小为[4, 10, 5]；
- 每个GPU分别进行前项计算过程；
- 前向过程计算完后，pytorch再从四个GPU中收集计算后的结果假设[4, 10, 5]，然后再按照次序将其拼接起来[16, 10, 5]，计算loss。
整个过程其实就是 同步模型参数→分别前向计算→计算损失→梯度反传
实现
定义模型
import torch
import torch.nn as nn
class Model(nn.Module):
def __init__(self):
self.linear = nn.Linear(5, 5)
def forward(self, inputs):
output = self.linear(inputs)
return output
model = Model()
定义优化器
optimizer = torch.optim.SGD(model.parameters(), lr = 0.01, momentum=0.9)
实现多GPU
# 假设就一个数据
data = torch.rand([16, 10, 5])
# 前向计算要求数据都放进GPU0里面
# device = torch.device('cuda:0')
# data = data.to(device)
data = data.cuda()
# 将网络同步到多个GPU中
model_p = torch.nn.DataParalle(model.cuda(), device_ids=[0, 1, 2, 3])
logits = model_p(inputs)
# 接下来计算loss
loss = crit(logits, target)
optimizer.zero_grad()
loss.backward()
optimizer.step()
注：模型要求所有的数据和初始网络被放置到GPU0，实际上并不需要，只需要保证数据和初始网路都在你所选择的多个gpu中的第一块上就行。
