GPU如何訓(xùn)練大批量模型？方法在這里

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

分布式計算

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

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

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

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

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

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

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

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

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

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

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

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

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

你要怎么做呢？

這個問題有一個簡單的解決方法：梯度累積。

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

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

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

以下是使用梯度累積訓(xùn)練模型的要點(diǎn)。在這個例子中，我們可以用一個大于 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

擴(kuò)展到極致

你可以在 GPU 上訓(xùn)練連一個樣本都無法加載的模型嗎？

如果你的架構(gòu)沒有太多跳過連接，這就是可能的！解決方案是使用梯度檢查點(diǎn)（gradient-checkpointing）來節(jié)省計算資源。