在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