我正在计算nn.Conv2d层内每对内核之间的累积距离。但是，对于较大的层，使用具有12gb内存的Titan X会耗尽内存。我想知道是否可以将这样的计算结果划分为两个GPU。 代码如下：
unique_ptr其中def ac_distance(layer):
total = 0
for p in layer.weight:
for q in layer.weight:
total += distance(p,q)
return total
是layer的实例，而distance返回p和q之差的总和。但是，我无法分离该图，因为稍后需要它。我尝试将模型包装在nn.DataParallel周围，但是nn.Conv2d中的所有计算仅使用1 gpu完成，但是使用这两者进行训练。
答案 0 :(得分：0)
并行训练神经网络可以通过两种方式实现。
正如您在问题中所提出的，您想将计算划分为第二类。没有开箱即用的方式来实现模型并行性。 PyTorch使用torch.distributed包为并行处理提供了原语。 tutorial全面介绍了该软件包的详细信息，您可以制定一种方法来实现所需的模型并行性。
但是，实现模型并行性可能非常复杂。一般的方法是使用torch.nn.DataParallel或torch.nn.DistributedDataParallel进行数据并行化。在这两种方法中，您将在两个不同的GPU上运行相同的模型，但是一大批将被分成两个较小的块。梯度将累积在单个GPU上并进行优化。通过使用多重处理，优化在Dataparallel中的单个GPU上进行，并且在DistributedDataParallel中的GPU上并行进行。
在您的情况下，如果使用DataParallel，则计算仍将在两个不同的GPU上进行。如果您发现GPU使用情况不平衡，可能是由于DataParallel的设计方式所致。您可以尝试使用DistributedDataParallel，这是根据docs在多个GPU上进行训练的最快方法。
还有其他方法可以处理非常大的批次。 article会详细介绍它们，我相信这会有所帮助。几个要点：