13.5.1。拆分問(wèn)題

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

首先，我們可以跨多個(gè) GPU 劃分網(wǎng)絡(luò)。也就是說(shuō)，每個(gè) GPU 將流入特定層的數(shù)據(jù)作為輸入，跨多個(gè)后續(xù)層處理數(shù)據(jù)，然后將數(shù)據(jù)發(fā)送到下一個(gè) GPU。與單個(gè) GPU 可以處理的數(shù)據(jù)相比，這使我們能夠使用更大的網(wǎng)絡(luò)處理數(shù)據(jù)。此外，可以很好地控制每個(gè) GPU 的內(nèi)存占用量（它只占網(wǎng)絡(luò)總占用量的一小部分）。

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

其次，我們可以分層拆分工作。例如，與其在單個(gè) GPU 上計(jì)算 64 個(gè)通道，不如將問(wèn)題拆分到 4 個(gè) GPU，每個(gè) GPU 生成 16 個(gè)通道的數(shù)據(jù)。同樣，對(duì)于全連接層，我們可以拆分輸出單元的數(shù)量。 圖 13.5.1（取自 Krizhevsky等人（2012 年））說(shuō)明了這種設(shè)計(jì)，其中這種策略用于處理內(nèi)存占用非常?。ó?dāng)時(shí)為 2 GB）的 GPU。如果通道（或單元）的數(shù)量不太小，這就可以在計(jì)算方面實(shí)現(xiàn)良好的縮放。此外，由于可用內(nèi)存線性擴(kuò)展，多個(gè) GPU 可以處理越來(lái)越大的網(wǎng)絡(luò)。

圖 13.5.1由于 GPU 內(nèi)存有限，原始 AlexNet 設(shè)計(jì)中的模型并行性。

然而，我們需要大量的同步或屏障操作，因?yàn)槊恳粚佣家蕾?lài)于所有其他層的結(jié)果。此外，需要傳輸?shù)臄?shù)據(jù)量可能比跨 GPU 分布層時(shí)更大。因此，由于帶寬成本和復(fù)雜性，我們不推薦這種方法。

最后，我們可以跨多個(gè) GPU 對(duì)數(shù)據(jù)進(jìn)行分區(qū)。這樣，所有 GPU 都執(zhí)行相同類(lèi)型的工作，盡管觀察結(jié)果不同。在每個(gè)小批量訓(xùn)練數(shù)據(jù)之后，梯度在 GPU 之間聚合。這是最簡(jiǎn)單的方法，適用于任何情況。我們只需要在每個(gè)小批量之后進(jìn)行同步。也就是說(shuō)，非常希望在其他仍在計(jì)算的同時(shí)開(kāi)始交換梯度參數(shù)。此外，更大數(shù)量的 GPU 會(huì)導(dǎo)致更大的小批量大小，從而提高訓(xùn)練效率。然而，添加更多 GPU 并不能讓我們訓(xùn)練更大的模型。

圖 13.5.2多 GPU 上的并行化。從左到右：原始問(wèn)題、網(wǎng)絡(luò)分區(qū)、分層分區(qū)、數(shù)據(jù)并行。

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

13.5.2。數(shù)據(jù)并行

假設(shè)有k機(jī)器上的 GPU。給定要訓(xùn)練的模型，每個(gè) GPU 將獨(dú)立維護(hù)一組完整的模型參數(shù)，盡管 GPU 之間的參數(shù)值是相同且同步的。例如，圖 13.5.3說(shuō)明了在以下情況下使用數(shù)據(jù)并行性進(jìn)行訓(xùn)練k=2.

圖 13.5.3在兩個(gè) GPU 上使用數(shù)據(jù)并行計(jì)算小批量隨機(jī)梯度下降。

一般來(lái)說(shuō)，訓(xùn)練過(guò)程如下：

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

• 每個(gè) GPU 根據(jù)分配給它的小批量子集計(jì)算模型參數(shù)的損失和梯度。

• 每個(gè)的局部梯度kGPU 被聚合以獲得當(dāng)前的小批量隨機(jī)梯度。

• 聚合梯度被重新分配給每個(gè) GPU。

• 每個(gè) GPU 使用這個(gè)小批量隨機(jī)梯度來(lái)更新它維護(hù)的完整模型參數(shù)集。

請(qǐng)注意，在實(shí)踐中我們增加了小批量大小k-訓(xùn)練時(shí)折疊kGPU 這樣每個(gè) GPU 都有相同數(shù)量的工作要做，就好像我們只在單個(gè) GPU 上訓(xùn)練一樣。在 16-GPU 服務(wù)器上，這會(huì)大大增加小批量大小，我們可能不得不相應(yīng)地增加學(xué)習(xí)率。另請(qǐng)注意，第 8.5 節(jié)中的批量歸一化需要進(jìn)行調(diào)整，例如，通過(guò)為每個(gè) GPU 保留一個(gè)單獨(dú)的批量歸一化系數(shù)。下面我們將使用玩具網(wǎng)絡(luò)來(lái)說(shuō)明多 GPU 訓(xùn)練。

%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。玩具網(wǎng)絡(luò)

我們使用7.6 節(jié)中介紹的 LeNet （稍作修改）。我們從頭開(kāi)始定義它以詳細(xì)說(shuō)明參數(shù)交換和同步。

# 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