GPU如何訓練大批量模型？方法在這里

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

　　分布式計算

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

　　但在多數情況下，隨機梯度下降算法需要很大批量才能得出不錯的結果。

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

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

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

　　我們將著(zhù)重探討以下問(wèn)題：

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

　　如何盡可能高效地利用多 GPU 機器;

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

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

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

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

　　但你很確定將批量加倍可以?xún)?yōu)化結果。

　　你要怎么做呢?

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

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

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

　　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() 運算期間，為每個(gè)參數計算梯度，并將其存儲在與每個(gè)參數相關(guān)聯(lián)的張量——parameter.grad 中。

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

　　以下是使用梯度累積訓練模型的要點(diǎn)。在這個(gè)例子中，我們可以用一個(gè)大于 GPU 最大容量的 accumulation_steps 批量進(jì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