在 PyTorch 训练之外使用多个 GPU答案

【问题标题】：Using Multiple GPUs outside of training in PyTorch在 PyTorch 训练之外使用多个 GPU
【发布时间】：2019-09-01 14:15:28
【问题描述】：

我正在计算 nn.Conv2d 层内每对内核之间的累积距离。然而，对于大型层，使用具有 12gb 内存的 Titan X 会耗尽内存。我想知道是否可以将此类计算划分为两个 GPU。代码如下：

def ac_distance(layer):
    total = 0
    for p in layer.weight:
      for q in layer.weight:
         total += distance(p,q)
    return total

其中layer 是nn.Conv2d 的实例，距离返回p 和q 之差的总和。但是，我无法分离图表，因为我稍后需要它。我尝试将我的模型包裹在 nn.DataParallel 周围，但 ac_distance 中的所有计算仅使用 1 个 gpu 完成，但它同时使用两者进行训练。

【问题讨论】：

标签： neural-network conv-neural-network pytorch

【解决方案1】：

训练神经网络时的并行性可以通过两种方式实现。

数据并行性 - 将大批量拆分为两个并执行相同的操作集，但分别在两个不同的 GPU 上进行
模型并行性 - 拆分计算并在不同的 GPU 上运行它们

正如您在问题中所问的，您想拆分属于第二类的计算。没有开箱即用的方法来实现模型并行性。 PyTorch 使用 torch.distributed 包提供用于并行处理的原语。这个tutorial 全面介绍了包的细节，您可以设计一种方法来实现您需要的模型并行性。

但是，实现模型并行性可能非常复杂。一般的方法是使用torch.nn.DataParallel 或torch.nn.DistributedDataParallel 进行数据并行。在这两种方法中，您都将在两个不同的 GPU 上运行相同的模型，但是一个大批量将被分成两个较小的块。梯度将在单个 GPU 上累积并进行优化。通过使用多处理在 Dataparallel 中的单个 GPU 上进行优化，并在 DistributedDataParallel 中跨 GPU 并行进行。

在您的情况下，如果您使用 DataParallel，计算仍将在两个不同的 GPU 上进行。如果您注意到 GPU 使用不平衡，可能是因为 DataParallel 的设计方式。您可以尝试使用DistributedDataParallel，这是根据docs 在多个 GPU 上训练的最快方法。

还有其他方法可以处理非常大的批次。这个article 详细介绍了它们，我相信它会有所帮助。几个重点：

对大批量进行梯度累积
使用 DataParallel
如果这还不够，请使用 DistributedDataParallel

【讨论】：