深度学习：Pytorch分布式训练-模型并行

模型并行主要利用to(device)函数将模型和数据（Tensor张量）放置在适当设备上，其余代码基本无需额外改动。
以下是一个简单的模型并行的代码示例：

import torch
import torch.nn as nn
import torch.optim as optim


class DemoModel(nn.Module):
    def __init__(self):
        super(DemoModel, self).__init__()
        self.net1 = torch.nn.Linear(10, 10).to('cuda:0')
        self.relu = torch.nn.ReLU()
        self.net2 = torch.nn.Linear(10, 5).to('cuda:1')

    def forward(self, x):
        x = self.relu(self.net1(x.to('cuda:0')))
        return self.net2(x.to('cuda:1'))

model = DemoModel()
loss_fn = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)

optimizer.zero_grad()
outputs = model(torch.randn(20, 10))
labels = torch.randn(20, 5).to('cuda:1')
loss_fn(outputs, labels).backward()
optimizer.step()

注意调用损失函数时，您只需要确保标签与输出位于同一设备上。不难看出，此模型并行的方法效率相对较低，因为在任何时间点，两个 GPU 中只有一个在工作，而另一个则处于闲置状态。而且中间过程变量从cuda:0复制到cuda:1，又会需要额外的开销。因此可以引入流水线并行来进行加速。

在以下代码示例中，采取将输入数据批次划分为 20 组。由于 PyTorch 异步启动 CUDA 操作，因此可以不需要生成多个线程来实现并发。值得注意的是，使用较小的结果split_size会导致许多微小的 CUDA 内核启动，而使用较大的split_size会导致在第一次和最后一次数据划分期间存在相对较长的空闲时间。因此split_size对于特定实验可能有一个最佳配置，可以多次尝试最佳的超参数。

class PipelineParallelResNet50(ModelParallelResNet50):
    def __init__(self, split_size=20, *args, **kwargs):
        super(PipelineParallelResNet50, self).__init__(*args, **kwargs)
        self.split_size = split_size

    def forward(self, x):
        splits = iter(x.split(self.split_size, dim=0))
        s_next = next(splits)
        s_prev = self.seq1(s_next).to('cuda:1')
        ret = []

        for s_next in splits:
            # A. ``s_prev`` runs on ``cuda:1``
            s_prev = self.seq2(s_prev)
            ret.append(self.fc(s_prev.view(s_prev.size(0), -1)))

            # B. ``s_next`` runs on ``cuda:0``, which can run concurrently with A
            s_prev = self.seq1(s_next).to('cuda:1')

        s_prev = self.seq2(s_prev)
        ret.append(self.fc(s_prev.view(s_prev.size(0), -1)))

        return torch.cat(ret)