深度學(xué)習(xí)模型和數(shù)據(jù)集的規(guī)模增長(zhǎng)速度已經(jīng)讓 GPU 算力也開始捉襟見肘，如果你的 GPU 連一個(gè)樣本都容不下，你要如何訓(xùn)練大批量模型?通過本文介紹的方法，我們可以在訓(xùn)練批量甚至單個(gè)訓(xùn)練樣本大于 GPU 內(nèi)存時(shí)，在單個(gè)或多個(gè) GPU 服務(wù)器上訓(xùn)練模型。

　　分布式計(jì)算

　　2018 年的大部分時(shí)間我都在試圖訓(xùn)練神經(jīng)網(wǎng)絡(luò)時(shí)克服 GPU 極限。無論是在含有 1.5 億個(gè)參數(shù)的語(yǔ)言模型(如 OpenAI 的大型生成預(yù)訓(xùn)練 Transformer 或最近類似的 BERT 模型)還是饋入 3000 萬個(gè)元素輸入的元學(xué)習(xí)神經(jīng)網(wǎng)絡(luò)(如我們?cè)谝黄?ICLR 論文《Meta-Learning a Dynamical Language Model》中提到的模型)，我都只能在 GPU 上處理很少的訓(xùn)練樣本。

　　但在多數(shù)情況下，隨機(jī)梯度下降算法需要很大批量才能得出不錯(cuò)的結(jié)果。

　　如果你的 GPU 只能處理很少的樣本，你要如何訓(xùn)練大批量模型?

　　有幾個(gè)工具、技巧可以幫助你解決上述問題。在本文中，我將自己用過、學(xué)過的東西整理出來供大家參考。

　　在這篇文章中，我將主要討論 PyTorch 框架。有部分工具尚未包括在 PyTorch(1.0 版本)中，因此我也寫了自定義代碼。

　　我們將著重探討以下問題：

　　在訓(xùn)練批量甚至單個(gè)訓(xùn)練樣本大于 GPU 內(nèi)存，要如何在單個(gè)或多個(gè) GPU 服務(wù)器上訓(xùn)練模型;

　　如何盡可能高效地利用多 GPU 機(jī)器;

　　在分布式設(shè)備上使用多個(gè)機(jī)器的最簡(jiǎn)單訓(xùn)練方法。

　　在一個(gè)或多個(gè) GPU 上訓(xùn)練大批量模型

　　你建的模型不錯(cuò)，在這個(gè)簡(jiǎn)潔的任務(wù)中可能成為新的 SOTA，但每次嘗試在一個(gè)批量處理更多樣本時(shí)，你都會(huì)得到一個(gè) CUDA RuntimeError：內(nèi)存不足。

　　這位網(wǎng)友指出了你的問題!

　　但你很確定將批量加倍可以優(yōu)化結(jié)果。

　　你要怎么做呢?

　　這個(gè)問題有一個(gè)簡(jiǎn)單的解決方法：梯度累積。

　　梯度下降優(yōu)化算法的五個(gè)步驟。

　　與之對(duì)等的 PyTorch 代碼也可以寫成以下五行：

　　predictions = model(inputs) # Forward pass

　　loss = loss_function(predictions, labels) # Compute loss function

　　loss.backward() # Backward pass

　　optimizer.step() # Optimizer step

　　predictions = model(inputs) # Forward pass with new parameters

　　在 loss.backward() 運(yùn)算期間，為每個(gè)參數(shù)計(jì)算梯度，并將其存儲(chǔ)在與每個(gè)參數(shù)相關(guān)聯(lián)的張量——parameter.grad 中。

　　累積梯度意味著，在調(diào)用 optimizer.step() 實(shí)施一步梯度下降之前，我們會(huì)對(duì) parameter.grad 張量中的幾個(gè)反向運(yùn)算的梯度求和。在 PyTorch 中這一點(diǎn)很容易實(shí)現(xiàn)，因?yàn)樘荻葟埩吭诓徽{(diào)用 model.zero_grad() 或 optimizer.zero_grad() 的情況下不會(huì)重置。如果損失在訓(xùn)練樣本上要取平均，我們還需要除以累積步驟的數(shù)量。

　　以下是使用梯度累積訓(xùn)練模型的要點(diǎn)。在這個(gè)例子中，我們可以用一個(gè)大于 GPU 最大容量的 accumulation_steps 批量進(jìn)行訓(xùn)練：

　　model.zero_grad() # Reset gradients tensors

　　for i, (inputs, labels) in enumerate(training_set):

　　predictions = model(inputs) # Forward pass

　　loss = loss_function(predictions, labels) # Compute loss function

　　loss = loss / accumulation_steps # Normalize our loss (if averaged)

　　loss.backward() # Backward pass

　　if (i+1) % accumulation_steps == 0: # Wait for several backward steps

　　optimizer.step() # Now we can do an optimizer step

　　model.zero_grad() # Reset gradients tensors

　　if (i+1) % evaluation_steps == 0: # Evaluate the model when we...

　　evaluate_model() # ...have no gradients accumulated