在PyTorch培训之外使用多个GPU

Question

我正在计算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实例，而distance返回p和q之差的总和。 但是，我无法分离该图，因为稍后需要它。 我尝试将模型包装在nn.DataParallel周围，但是ac_distance中的所有计算仅使用1 gpu完成，但是使用这两者进行训练。

Answer 1

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

1. 数据并行化-将大批分成两部分，并执行相同的操作集，但分别在两个不同的GPU上进行
2. 并行模型-拆分计算并在不同的GPU上运行

正如您在问题中所提出的，您想将计算划分为第二类。 没有开箱即用的方式来实现模型并行性。 PyTorch提供了使用torch.distributed包进行并行处理的原语。 本教程全面介绍了该软件包的详细信息，您可以制定一种方法来实现所需的模型并行性。

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

在您的情况下，如果使用DataParallel ，则计算仍将在两个不同的GPU上进行。 如果您发现GPU使用情况不平衡，可能是由于DataParallel设计方式DataParallel 。 根据文档，您可以尝试使用DistributedDataParallel ，这是在多个GPU上进行训练的最快方法。

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

• 对较大的批次进行梯度累积
• 使用DataParallel
• 如果那还不够，请使用DistributedDataParallel