PyTorch教程13.3之自動(dòng)并行

13.3.1。GPU 上的并行計(jì)算

讓我們首先定義一個(gè)要測(cè)試的參考工作負(fù)載：run 下面的函數(shù)使用分配到兩個(gè)變量中的數(shù)據(jù)在我們選擇的設(shè)備上執(zhí)行 10 次矩陣-矩陣乘法：x_gpu1和 x_gpu2。

devices = d2l.try_all_gpus()
def run(x):
  return [x.mm(x) for _ in range(50)]

x_gpu1 = torch.rand(size=(4000, 4000), device=devices[0])
x_gpu2 = torch.rand(size=(4000, 4000), device=devices[1])

run(x_gpu1)
run(x_gpu2) # Warm-up all devices
torch.cuda.synchronize(devices[0])
torch.cuda.synchronize(devices[1])

with d2l.Benchmark('GPU1 time'):
  run(x_gpu1)
  torch.cuda.synchronize(devices[0])

with d2l.Benchmark('GPU2 time'):
  run(x_gpu2)
  torch.cuda.synchronize(devices[1])

GPU1 time: 0.4967 sec
GPU2 time: 0.5151 sec

with d2l.Benchmark('GPU1 & GPU2'):
  run(x_gpu1)
  run(x_gpu2)
  torch.cuda.synchronize()

GPU1 & GPU2: 0.5000 sec

devices = d2l.try_all_gpus()
def run(x):
  return [x.dot(x) for _ in range(50)]

x_gpu1 = np.random.uniform(size=(4000, 4000), ctx=devices[0])
x_gpu2 = np.random.uniform(size=(4000, 4000), ctx=devices[1])

run(x_gpu1) # Warm-up both devices
run(x_gpu2)
npx.waitall()

with d2l.Benchmark('GPU1 time'):
  run(x_gpu1)
  npx.waitall()

with d2l.Benchmark('GPU2 time'):
  run(x_gpu2)
  npx.waitall()

GPU1 time: 0.5233 sec
GPU2 time: 0.5158 sec

with d2l.Benchmark('GPU1 & GPU2'):
  run(x_gpu1)
  run(x_gpu2)
  npx.waitall()

GPU1 & GPU2: 0.5214 sec

在上述情況下，總執(zhí)行時(shí)間小于其各部分的總和，因?yàn)樯疃葘W(xué)習(xí)框架會(huì)自動(dòng)安排兩個(gè) GPU 設(shè)備上的計(jì)算，而不需要代表用戶(hù)編寫(xiě)復(fù)雜的代碼。

13.3.2。并行計(jì)算與通信

在許多情況下，我們需要在不同設(shè)備之間移動(dòng)數(shù)據(jù)，比如在 CPU 和 GPU 之間，或者在不同 GPU 之間。例如，當(dāng)我們想要執(zhí)行分布式優(yōu)化時(shí)會(huì)發(fā)生這種情況，我們需要在多個(gè)加速器卡上聚合梯度。讓我們通過(guò)在 GPU 上計(jì)算然后將結(jié)果復(fù)制回 CPU 來(lái)對(duì)此進(jìn)行模擬。

def copy_to_cpu(x, non_blocking=False):
  return [y.to('cpu', non_blocking=non_blocking) for y in x]

with d2l.Benchmark('Run on GPU1'):
  y = run(x_gpu1)
  torch.cuda.synchronize()

with d2l.Benchmark('Copy to CPU'):
  y_cpu = copy_to_cpu(y)
  torch.cuda.synchronize()

Run on GPU1: 0.5019 sec
Copy to CPU: 2.7168 sec

with d2l.Benchmark('Run on GPU1 and copy to CPU'):
  y = run(x_gpu1)
  y_cpu = copy_to_cpu(y, True)
  torch.cuda.synchronize()

Run on GPU1 and copy to CPU: 2.4682 sec

def copy_to_cpu(x):
  return [y.copyto(npx.cpu()) for y in x]

with d2l.Benchmark('Run on GPU1'):
  y = run(x_gpu1)
  npx.waitall()

with d2l.Benchmark('Copy to CPU'):
  y_cpu = copy_to_cpu(y)
  npx.waitall()

Run on GPU1: 0.5796 sec
Copy to CPU: 3.0989 sec

with d2l.Benchmark('Run on GPU1 and copy to CPU'):
  y = run(x_gpu1)
  y_cpu = copy_to_cpu(y)
  npx.waitall()

Run on GPU1 and copy to CPU: 3.3488 sec

兩個(gè)操作所需的總時(shí)間（正如預(yù)期的那樣）小于它們各部分的總和。請(qǐng)注意，此任務(wù)不同于并行計(jì)算，因?yàn)樗褂貌煌馁Y源：CPU 和 GPU 之間的總線(xiàn)。事實(shí)上，我們可以同時(shí)在兩個(gè)設(shè)備上進(jìn)行計(jì)算和通信。如上所述，計(jì)算和通信之間存在依賴(lài)關(guān)系：y[i]必須在將其復(fù)制到 CPU 之前進(jìn)行計(jì)算。幸運(yùn)的是，系統(tǒng)可以y[i-1]邊計(jì)算邊 復(fù)制y[i]，以減少總運(yùn)行時(shí)間。

我們以在一個(gè) CPU 和兩個(gè) GPU 上進(jìn)行訓(xùn)練時(shí)簡(jiǎn)單的兩層 MLP 的計(jì)算圖及其依賴(lài)關(guān)系的圖示作為結(jié)尾，如圖13.3.1所示。手動(dòng)安排由此產(chǎn)生的并行程序?qū)⒎浅Ｍ纯唷?/font>這就是擁有基于圖形的計(jì)算后端進(jìn)行優(yōu)化的優(yōu)勢(shì)所在。

圖 13.3.1兩層 MLP 在一個(gè) CPU 和兩個(gè) GPU 上的計(jì)算圖及其依賴(lài)關(guān)系。