PyTorch教程-13.5。在多個 GPU 上進(jìn)行訓(xùn)練

到目前為止，我們討論了如何在 CPU 和 GPU 上高效地訓(xùn)練模型。在13.3 節(jié)中，我們甚至展示了深度學(xué)習(xí)框架如何允許人們在它們之間自動并行計算和通信 。我們還在6.7 節(jié)中展示了如何使用nvidia-smi 命令列出計算機(jī)上所有可用的 GPU。我們沒有討論的是如何真正并行化深度學(xué)習(xí)訓(xùn)練。相反，我們暗示傳遞一個會以某種方式將數(shù)據(jù)拆分到多個設(shè)備并使其工作。本節(jié)填寫詳細(xì)信息并展示如何從頭開始并行訓(xùn)練網(wǎng)絡(luò)。有關(guān)如何利用高級 API 中的功能的詳細(xì)信息歸入 第 13.6 節(jié). 我們假設(shè)您熟悉 minibatch 隨機(jī)梯度下降算法，例如12.5 節(jié)中描述的算法。

13.5.1。拆分問題

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

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

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

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

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

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

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

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

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

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

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

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

一般來說，訓(xùn)練過程如下：

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

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

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

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

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

請注意，在實(shí)踐中我們增加了小批量大小k-訓(xùn)練時折疊kGPU 這樣每個 GPU 都有相同數(shù)量的工作要做，就好像我們只在單個 GPU 上訓(xùn)練一樣。在 16-GPU 服務(wù)器上，這會大大增加小批量大小，我們可能不得不相應(yīng)地增加學(xué)習(xí)率。另請注意，第 8.5 節(jié)中的批量歸一化需要進(jìn)行調(diào)整，例如，通過為每個 GPU 保留一個單獨(dú)的批量歸一化系數(shù)。下面我們將使用玩具網(wǎng)絡(luò)來說明多 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 （稍作修改）。我們從頭開始定義它以詳細(xì)說明參數(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.shape[0], -1)
  h3_linear = np.dot(h2, params[4]) + params[5]
  h3 = npx.relu(h3_linear)
  y_hat = np.dot(h3, params[6]) + params[7]
  return y_hat

# Cross-entropy loss function
loss = gluon.loss.SoftmaxCrossEntropyLoss()

13.5.4。數(shù)據(jù)同步

為了進(jìn)行高效的多 GPU 訓(xùn)練，我們需要兩個基本操作。首先，我們需要能夠?qū)?shù)列表分發(fā)到多個設(shè)備并附加梯度 ( get_params)。沒有參數(shù)就不可能在 GPU 上評估網(wǎng)絡(luò)。其次，我們需要能夠跨多個設(shè)備對參數(shù)求和，即，我們需要一個 allreduce函數(shù)。

def get_params(params, device):
  new_params = [p.to(device) for p in params]
  for p in new_params:
    p.requires_grad_()
  return new_params

def get_params(params, device):
  new_params = [p.copyto(device) for p in params]
  for p in new_params:
    p.attach_grad()
  return new_params

讓我們通過將模型參數(shù)復(fù)制到一個 GPU 來嘗試一下。

new_params = get_params(params, d2l.try_gpu(0))
print('b1 weight:', new_params[1])
print('b1 grad:', new_params[1].grad)

b1 weight: tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
    device='cuda:0', requires_grad=True)
b1 grad: None

new_params = get_params(params, d2l.try_gpu(0))
print('b1 weight:', new_params[1])
print('b1 grad:', new_params[1].grad)

b1 weight: [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.] @gpu(0)
b1 grad: [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.] @gpu(0)

由于我們還沒有執(zhí)行任何計算，因此關(guān)于偏置參數(shù)的梯度仍然為零。現(xiàn)在讓我們假設(shè)我們有一個分布在多個 GPU 上的向量。以下allreduce 函數(shù)將所有向量相加并將結(jié)果廣播回所有 GPU。請注意，為了使其正常工作，我們需要將數(shù)據(jù)復(fù)制到設(shè)備以累積結(jié)果。

def allreduce(data):
  for i in range(1, len(data)):
    data[0][:] += data[i].to(data[0].device)
  for i in range(1, len(data)):
    data[i][:] = data[0].to(data[i].device)

def allreduce(data):
  for i in range(1, len(data)):
    data[0][:] += data[i].copyto(data[0].ctx)
  for i in range(1, len(data)):
    data[0].copyto(data[i])

讓我們通過在不同設(shè)備上創(chuàng)建具有不同值的向量并聚合它們來對此進(jìn)行測試。

data = [torch.ones((1, 2), device=d2l.try_gpu(i)) * (i + 1) for i in range(2)]
print('before allreduce:n', data[0], 'n', data[1])
allreduce(data)
print('after allreduce:n', data[0], 'n', data[1])

before allreduce:
 tensor([[1., 1.]], device='cuda:0')
 tensor([[2., 2.]], device='cuda:1')
after allreduce:
 tensor([[3., 3.]], device='cuda:0')
 tensor([[3., 3.]], device='cuda:1')

data = [np.ones((1, 2), ctx=d2l.try_gpu(i)) * (i + 1) for i in range(2)]
print('before allreduce:n', data[0], 'n', data[1])
allreduce(data)
print('after allreduce:n', data[0], 'n', data[1])

before allreduce:
 [[1. 1.]] @gpu(0)
 [[2. 2.]] @gpu(1)
after allreduce:
 [[3. 3.]] @gpu(0)
 [[3. 3.]] @gpu(1)

13.5.5。分發(fā)數(shù)據(jù)

我們需要一個簡單的效用函數(shù)來在多個 GPU 上平均分配一個小批量。例如，在兩個 GPU 上，我們希望將一半數(shù)據(jù)復(fù)制到其中一個 GPU。由于它更方便、更簡潔，我們使用深度學(xué)習(xí)框架的內(nèi)置函數(shù)來嘗試一下4×5矩陣。

data = torch.arange(20).reshape(4, 5)
devices = [torch.device('cuda:0'), torch.device('cuda:1')]
split = nn.parallel.scatter(data, devices)
print('input :', data)
print('load into', devices)
print('output:', split)

input : tensor([[ 0, 1, 2, 3, 4],
    [ 5, 6, 7, 8, 9],
    [10, 11, 12, 13, 14],
    [15, 16, 17, 18, 19]])
load into [device(type='cuda', index=0), device(type='cuda', index=1)]
output: (tensor([[0, 1, 2, 3, 4],
    [5, 6, 7, 8, 9]], device='cuda:0'), tensor([[10, 11, 12, 13, 14],
    [15, 16, 17, 18, 19]], device='cuda:1'))

data = np.arange(20).reshape(4, 5)
devices = [npx.gpu(0), npx.gpu(1)]
split = gluon.utils.split_and_load(data, devices)
print('input :', data)
print('load into', devices)
print('output:', split)

input : [[ 0. 1. 2. 3. 4.]
 [ 5. 6. 7. 8. 9.]
 [10. 11. 12. 13. 14.]
 [15. 16. 17. 18. 19.]]
load into [gpu(0), gpu(1)]
output: [array([[0., 1., 2., 3., 4.],
    [5., 6., 7., 8., 9.]], ctx=gpu(0)), array([[10., 11., 12., 13., 14.],
    [15., 16., 17., 18., 19.]], ctx=gpu(1))]

為了以后重用，我們定義了一個split_batch函數(shù)來拆分?jǐn)?shù)據(jù)和標(biāo)簽。

#@save
def split_batch(X, y, devices):
  """Split `X` and `y` into multiple devices."""
  assert X.shape[0] == y.shape[0]
  return (nn.parallel.scatter(X, devices),
      nn.parallel.scatter(y, devices))

#@save
def split_batch(X, y, devices):
  """Split `X` and `y` into multiple devices."""
  assert X.shape[0] == y.shape[0]
  return (gluon.utils.split_and_load(X, devices),
      gluon.utils.split_and_load(y, devices))

13.5.6。訓(xùn)練

現(xiàn)在我們可以在單個 minibatch 上實(shí)現(xiàn)多 GPU 訓(xùn)練。它的實(shí)現(xiàn)主要基于本節(jié)中描述的數(shù)據(jù)并行方法。我們將使用我們剛剛討論的輔助函數(shù)allreduce和split_and_load來同步多個 GPU 之間的數(shù)據(jù)。請注意，我們不需要編寫任何特定代碼來實(shí)現(xiàn)并行性。由于計算圖在 minibatch 中沒有任何跨設(shè)備依賴性，因此它會自動并行執(zhí)行。

def train_batch(X, y, device_params, devices, lr):
  X_shards, y_shards = split_batch(X, y, devices)
  # Loss is calculated separately on each GPU
  ls = [loss(lenet(X_shard, device_W), y_shard).sum()
     for X_shard, y_shard, device_W in zip(
       X_shards, y_shards, device_params)]
  for l in ls: # Backpropagation is performed separately on each GPU
    l.backward()
  # Sum all gradients from each GPU and broadcast them to all GPUs
  with torch.no_grad():
    for i in range(len(device_params[0])):
      allreduce([device_params[c][i].grad for c in range(len(devices))])
  # The model parameters are updated separately on each GPU
  for param in device_params:
    d2l.sgd(param, lr, X.shape[0]) # Here, we use a full-size batch

def train_batch(X, y, device_params, devices, lr):
  X_shards, y_shards = split_batch(X, y, devices)
  with autograd.record(): # Loss is calculated separately on each GPU
    ls = [loss(lenet(X_shard, device_W), y_shard)
       for X_shard, y_shard, device_W in zip(
         X_shards, y_shards, device_params)]
  for l in ls: # Backpropagation is performed separately on each GPU
    l.backward()
  # Sum all gradients from each GPU and broadcast them to all GPUs
  for i in range(len(device_params[0])):
    allreduce([device_params[c][i].grad for c in range(len(devices))])
  # The model parameters are updated separately on each GPU
  for param in device_params:
    d2l.sgd(param, lr, X.shape[0]) # Here, we use a full-size batch

現(xiàn)在，我們可以定義訓(xùn)練函數(shù)。它與前面章節(jié)中使用的略有不同：我們需要分配 GPU 并將所有模型參數(shù)復(fù)制到所有設(shè)備。顯然，每個批次都是使用train_batch處理多個 GPU 的函數(shù)進(jìn)行處理的。為了方便（和代碼的簡潔性），我們在單個 GPU 上計算精度，但由于其他 GPU 處于空閑狀態(tài)，因此效率很低。

def train(num_gpus, batch_size, lr):
  train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
  devices = [d2l.try_gpu(i) for i in range(num_gpus)]
  # Copy model parameters to `num_gpus` GPUs
  device_params = [get_params(params, d) for d in devices]
  num_epochs = 10
  animator = d2l.Animator('epoch', 'test acc', xlim=[1, num_epochs])
  timer = d2l.Timer()
  for epoch in range(num_epochs):
    timer.start()
    for X, y in train_iter:
      # Perform multi-GPU training for a single minibatch
      train_batch(X, y, device_params, devices, lr)
      torch.cuda.synchronize()
    timer.stop()
    # Evaluate the model on GPU 0
    animator.add(epoch + 1, (d2l.evaluate_accuracy_gpu(
      lambda x: lenet(x, device_params[0]), test_iter, devices[0]),))
  print(f'test acc: {animator.Y[0][-1]:.2f}, {timer.avg():.1f} sec/epoch '
     f'on {str(devices)}')

def train(num_gpus, batch_size, lr):
  train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
  devices = [d2l.try_gpu(i) for i in range(num_gpus)]
  # Copy model parameters to `num_gpus` GPUs
  device_params = [get_params(params, d) for d in devices]
  num_epochs = 10
  animator = d2l.Animator('epoch', 'test acc', xlim=[1, num_epochs])
  timer = d2l.Timer()
  for epoch in range(num_epochs):
    timer.start()
    for X, y in train_iter:
      # Perform multi-GPU training for a single minibatch
      train_batch(X, y, device_params, devices, lr)
      npx.waitall()
    timer.stop()
    # Evaluate the model on GPU 0
    animator.add(epoch + 1, (d2l.evaluate_accuracy_gpu(
      lambda x: lenet(x, device_params[0]), test_iter, devices[0]),))
  print(f'test acc: {animator.Y[0][-1]:.2f}, {timer.avg():.1f} sec/epoch '
     f'on {str(devices)}')

讓我們看看它在單個 GPU 上的表現(xiàn)如何。我們首先使用 256 的批量大小和 0.2 的學(xué)習(xí)率。

train(num_gpus=1, batch_size=256, lr=0.2)

test acc: 0.84, 4.1 sec/epoch on [device(type='cuda', index=0)]

train(num_gpus=1, batch_size=256, lr=0.2)

test acc: 0.83, 6.3 sec/epoch on [gpu(0)]

通過保持 batch size 和 learning rate 不變，并將 GPU 的數(shù)量增加到 2，我們可以看到測試準(zhǔn)確率與之前的實(shí)驗(yàn)相比大致保持不變。就優(yōu)化算法而言，它們是相同的。不幸的是，這里沒有獲得有意義的加速：模型太小了；此外，我們只有一個小數(shù)據(jù)集，其中我們實(shí)現(xiàn)多 GPU 訓(xùn)練的稍微簡單的方法遭受了顯著的 Python 開銷。我們將遇到更復(fù)雜的模型和更復(fù)雜的并行化方式。讓我們看看 Fashion-MNIST 會發(fā)生什么。

train(num_gpus=2, batch_size=256, lr=0.2)

test acc: 0.83, 4.6 sec/epoch on [device(type='cuda', index=0), device(type='cuda', index=1)]

train(num_gpus=2, batch_size=256, lr=0.2)

test acc: 0.85, 13.7 sec/epoch on [gpu(0), gpu(1)]

13.5.7。概括

有多種方法可以在多個 GPU 上拆分深度網(wǎng)絡(luò)訓(xùn)練。我們可以在層之間、跨層或跨數(shù)據(jù)拆分它們。前兩者需要精心編排的數(shù)據(jù)傳輸。數(shù)據(jù)并行是最簡單的策略。

數(shù)據(jù)并行訓(xùn)練很簡單。但是，它增加了有效的小批量大小以提高效率。

在數(shù)據(jù)并行中，數(shù)據(jù)被拆分到多個 GPU，其中每個 GPU 執(zhí)行自己的前向和后向操作，隨后聚合梯度并將結(jié)果廣播回 GPU。

對于較大的小批量，我們可能會使用稍微增加的學(xué)習(xí)率。

13.5.8。練習(xí)

訓(xùn)練時kGPU，將小批量大小從 b到k?b，即，按 GPU 的數(shù)量進(jìn)行擴(kuò)展。

比較不同學(xué)習(xí)率的準(zhǔn)確性。它如何隨 GPU 數(shù)量擴(kuò)展？

實(shí)現(xiàn)一個更高效的allreduce函數(shù)，在不同的 GPU 上聚合不同的參數(shù)？為什么效率更高？

實(shí)現(xiàn)多 GPU 測試精度計算。