PyTorch教程13.5之在多個GPU上進行訓練

13.5.1。拆分問題

讓我們從一個簡單的計算機視覺問題和一個稍微陳舊的網絡開始，例如，具有多層卷積、池化，最后可能還有一些完全連接的層。也就是說，讓我們從一個看起來與 LeNet （LeCun等人，1998 年）或 AlexNet （Krizhevsky等人，2012 年）非常相似的網絡開始。給定多個 GPU（如果是桌面服務器則為 2 個，在 AWS g4dn.12xlarge 實例上為 4 個，在 p3.16xlarge 上為 8 個，或在 p2.16xlarge 上為 16 個），我們希望以實現良好加速的方式對訓練進行分區同時受益于簡單且可重現的設計選擇。畢竟，多個 GPU 會增加內存和計算能力。簡而言之，給定我們要分類的小批量訓練數據，我們有以下選擇。

首先，我們可以跨多個 GPU 劃分網絡。也就是說，每個 GPU 將流入特定層的數據作為輸入，跨多個后續層處理數據，然后將數據發送到下一個 GPU。與單個 GPU 可以處理的數據相比，這使我們能夠使用更大的網絡處理數據。此外，可以很好地控制每個 GPU 的內存占用量（它只占網絡總占用量的一小部分）。

然而，層（以及 GPU）之間的接口需要緊密同步。這可能很棘手，特別是如果層與層之間的計算工作負載沒有正確匹配。對于大量 GPU，問題會更加嚴重。層與層之間的接口也需要大量的數據傳輸，例如激活和梯度。這可能會超出 GPU 總線的帶寬。此外，計算密集型但順序的操作對于分區來說并不重要。參見例如Mirhoseini等人。( 2017 年)在這方面盡最大努力。這仍然是一個難題，尚不清楚是否有可能在非平凡問題上實現良好的（線性）縮放。我們不推薦它，除非有出色的框架或操作系統支持將多個 GPU 鏈接在一起。

其次，我們可以分層拆分工作。例如，與其在單個 GPU 上計算 64 個通道，不如將問題拆分到 4 個 GPU，每個 GPU 生成 16 個通道的數據。同樣，對于全連接層，我們可以拆分輸出單元的數量。 圖 13.5.1（取自 Krizhevsky等人（2012 年））說明了這種設計，其中這種策略用于處理內存占用非常?。ó敃r為 2 GB）的 GPU。如果通道（或單元）的數量不太小，這就可以在計算方面實現良好的縮放。此外，由于可用內存線性擴展，多個 GPU 可以處理越來越大的網絡。

圖 13.5.1由于 GPU 內存有限，原始 AlexNet 設計中的模型并行性。

然而，我們需要大量的同步或屏障操作，因為每一層都依賴于所有其他層的結果。此外，需要傳輸的數據量可能比跨 GPU 分布層時更大。因此，由于帶寬成本和復雜性，我們不推薦這種方法。

最后，我們可以跨多個 GPU 對數據進行分區。這樣，所有 GPU 都執行相同類型的工作，盡管觀察結果不同。在每個小批量訓練數據之后，梯度在 GPU 之間聚合。這是最簡單的方法，適用于任何情況。我們只需要在每個小批量之后進行同步。也就是說，非常希望在其他仍在計算的同時開始交換梯度參數。此外，更大數量的 GPU 會導致更大的小批量大小，從而提高訓練效率。然而，添加更多 GPU 并不能讓我們訓練更大的模型。

圖 13.5.2多 GPU 上的并行化。從左到右：原始問題、網絡分區、分層分區、數據并行。

圖 13.5.2描繪了多 GPU 上不同并行化方式的比較。總的來說，數據并行是最方便的方法，前提是我們可以訪問具有足夠大內存的 GPU。另請參閱 ( Li et al. , 2014 )以了解分布式訓練分區的詳細描述。在深度學習的早期，GPU 內存曾經是一個問題。到目前為止，除了最不尋常的情況外，所有問題都已解決。下面我們重點介紹數據并行性。

13.5.2。數據并行

假設有k機器上的 GPU。給定要訓練的模型，每個 GPU 將獨立維護一組完整的模型參數，盡管 GPU 之間的參數值是相同且同步的。例如，圖 13.5.3說明了在以下情況下使用數據并行性進行訓練k=2.

圖 13.5.3在兩個 GPU 上使用數據并行計算小批量隨機梯度下降。

一般來說，訓練過程如下：

• 在訓練的任何迭代中，給定一個隨機小批量，我們將批量中的示例分成k部分并將它們均勻地分布在 GPU 上。

• 每個 GPU 根據分配給它的小批量子集計算模型參數的損失和梯度。

• 每個的局部梯度kGPU 被聚合以獲得當前的小批量隨機梯度。

• 聚合梯度被重新分配給每個 GPU。

• 每個 GPU 使用這個小批量隨機梯度來更新它維護的完整模型參數集。

請注意，在實踐中我們增加了小批量大小k-訓練時折疊kGPU 這樣每個 GPU 都有相同數量的工作要做，就好像我們只在單個 GPU 上訓練一樣。在 16-GPU 服務器上，這會大大增加小批量大小，我們可能不得不相應地增加學習率。另請注意，第 8.5 節中的批量歸一化需要進行調整，例如，通過為每個 GPU 保留一個單獨的批量歸一化系數。下面我們將使用玩具網絡來說明多 GPU 訓練。

%matplotlib inline
import torch
from torch import nn
from torch.nn import functional as F
from d2l import torch as d2l

%matplotlib inline
from mxnet import autograd, gluon, np, npx
from d2l import mxnet as d2l

npx.set_np()

13.5.3。玩具網絡

我們使用7.6 節中介紹的 LeNet （稍作修改）。我們從頭開始定義它以詳細說明參數交換和同步。

# Initialize model parameters
scale = 0.01
W1 = torch.randn(size=(20, 1, 3, 3)) * scale
b1 = torch.zeros(20)
W2 = torch.randn(size=(50, 20, 5, 5)) * scale
b2 = torch.zeros(50)
W3 = torch.randn(size=(800, 128)) * scale
b3 = torch.zeros(128)
W4 = torch.randn(size=(128, 10)) * scale
b4 = torch.zeros(10)
params = [W1, b1, W2, b2, W3, b3, W4, b4]

# Define the model
def lenet(X, params):
  h1_conv = F.conv2d(input=X, weight=params[0], bias=params[1])
  h1_activation = F.relu(h1_conv)
  h1 = F.avg_pool2d(input=h1_activation, kernel_size=(2, 2), stride=(2, 2))
  h2_conv = F.conv2d(input=h1, weight=params[2], bias=params[3])
  h2_activation = F.relu(h2_conv)
  h2 = F.avg_pool2d(input=h2_activation, kernel_size=(2, 2), stride=(2, 2))
  h2 = h2.reshape(h2.shape[0], -1)
  h3_linear = torch.mm(h2, params[4]) + params[5]
  h3 = F.relu(h3_linear)
  y_hat = torch.mm(h3, params[6]) + params[7]
  return y_hat

# Cross-entropy loss function
loss = nn.CrossEntropyLoss(reduction='none')

# Initialize model parameters
scale = 0.01
W1 = np.random.normal(scale=scale, size=(20, 1, 3, 3))
b1 = np.zeros(20)
W2 = np.random.normal(scale=scale, size=(50, 20, 5, 5))
b2 = np.zeros(50)
W3 = np.random.normal(scale=scale, size=(800, 128))
b3 = np.zeros(128)
W4 = np.random.normal(scale=scale, size=(128, 10))
b4 = np.zeros(10)
params = [W1, b1, W2, b2, W3, b3, W4, b4]

# Define the model
def lenet(X, params):
  h1_conv = npx.convolution(data=X, weight=params[0], bias=params[1],
               kernel=(3, 3), num_filter=20)
  h1_activation = npx.relu(h1_conv)
  h1 = npx.pooling(data=h1_activation, pool_type='avg', kernel=(2, 2),
           stride=(2, 2))
  h2_conv = npx.convolution(data=h1, weight=params[2], bias=params[3],
               kernel=(5, 5), num_filter=50)
  h2_activation = npx.relu(h2_conv)
  h2 = npx.pooling(data=h2_activation, pool_type='avg', kernel=(2, 2),
           stride=(2, 2))
  h2 = h2.reshape(h2