Pytorch Multi-GPU原理与实现（单机多卡）

最近训练模型时候想要使用使用多GPU运算来提高计算速度，参考一些博客以及自己的动手实验搞懂了Pytorch的Multi-GPU原理。现在稍微整理一下。

原理

通常情况下，多GPU运算分为单机多卡和多机多卡，两者在pytorch上面的实现并不相同，因为多机时，需要多个机器之间的通信协议等设置。

pytorch实现单机多卡十分容易，其基本原理就是：加入我们一次性读入一个batch的数据, 其大小为[16, 10, 5]，我们有四张卡可以使用。那么计算过程遵循以下步骤：

1. 假设我们有4个GPU可以用，pytorch先把模型同步放到4个GPU中。
2. 那么首先将数据分为4份，按照次序放置到四个GPU的模型中，每一份大小为[4, 10, 5]；
3. 每个GPU分别进行前项计算过程；
4. 前向过程计算完后，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中的第一块上就行。