本文共 3112 字，大约阅读时间需要 10 分钟。
从PyTorch 0.2版本开始，PyTorch新增分布式GPU支持。
注意分布式和并行的区别：分布式是指有多个GPU在多台服务器上，而并行一般指的是一台服务器上的多个GPU。分布式涉及了服务器之间的通信，因此比较复杂，PyTorch封装了相应的接口，可以用几句简单的代码实现分布式训练。
分布式对普通用户来说比较遥远，因为搭建一个分布式集群的代价很大，使用也比较复杂。相比之下，一机多卡更现实。
如果服务器具有多个GPU，tensor.cuda（）方法会将tensor保存到第一块GPU上，等价于tensor.cuda（0）。此时如果想使用第二块GPU，需手动指定tensor.cuda（1）
Pytorch中多GPU并行计算
我们主要谈论的是单主机多GPU训练
- DataParallel实现简单，代码量较少，启动速度快一点。但速度较慢，且存在负载不均衡的问题。
- DistributedDataParallel本身是实现多机多卡的，但单机多卡也可以使用，配置稍复杂
DataParallel是Pytorch官方很久之前给的一种方案。单进程，多线程
DistributedDataParallel是更新一代的多卡训练方法。多进程
现在都是用DistributedDataParallel
注意 .to(device)就是把数据从内存放到GPU显存
Data Parallel
Model Parallel
Part of the model on CPU and part on the GPU
class DistributedModel(nn.Module): def __init__(self): super().__init__( embedding=nn.Embedding(1000, 10), rnn=nn.Linear(10, 10).to(device), ) def forward(self, x): # Compute embedding on CPU x = self.embedding(x) # Transfer to GPU x = x.to(device) # Compute RNN on GPU x = self.rnn(x) return x
Basic Usage
首先从包含两个线性层的toy model开始。 要在两个GPU上运行此模型，只需将每个线性层放在不同的GPU上，然后移动输入和中间输出以匹配设备。
import torchimport torch.nn as nnimport torch.optim as optimclass ToyModel(nn.Module): def __init__(self): super(ToyModel, 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 = ToyModel()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()除了5个 to(device) 调用将线性层和张量放置在适当的设备上之外，上面的ToyModel看起来非常类似于在单个GPU上实现它的方式。 to(device)是模型中唯一需要更改的地方。多GPU训练时 backward() 和torch.optim将自动处理梯度，就和模型在一个GPU上一样。 只需要注意在调用损失函数时，确保label与output位于同一设备上。
另一个例子
from torchvision.models.resnet import ResNet, Bottlenecknum_classes = 1000class ModelParallelResNet50(ResNet): def __init__(self, *args, **kwargs): super(ModelParallelResNet50, self).__init__( Bottleneck, [3, 4, 6, 3], num_classes=num_classes, *args, **kwargs) self.seq1 = nn.Sequential( self.conv1, self.bn1, self.relu, self.maxpool, self.layer1, self.layer2 ).to('cuda:0') self.seq2 = nn.Sequential( self.layer3, self.layer4, self.avgpool, ).to('cuda:1') self.fc.to('cuda:1') def forward(self, x): x = self.seq2(self.seq1(x).to('cuda:1')) return self.fc(x.view(x.size(0), -1))这就显示了如何将torchvision.models.resnet50（）分解为两个GPU。 这个想法是继承现有的ResNet模块，并在构建过程中将层拆分为两个GPU。 然后，通过相应地移动中间输出，覆盖前向方法来连接两个子模块。
对于模型太大而无法放入单个GPU的情况，上述实现解决了该问题。 但是要注意到，它将比在单个GPU上运行要慢。 这是因为在任何时间点，两个GPU中只有一个在工作，而另一个在那儿什么也没做。 由于中间输出需要在第2层和第3层之间从cuda:0复制到cuda:1，因此性能进一步下降
我们可以对此改进，因为我们知道两个GPU之一在整个执行过程中都处于闲置状态。 一种选择是将每个批次进一步划分为拆分流水线，以便当一个拆分到达第二个卡时，可以将下一个拆分馈入第一张卡。 这样，两个连续的拆分可以在两个GPU上同时运行。
Speed Up by Pipelining Inputs
转载地址：http://eqygf.baihongyu.com/