深度學(xué)習(xí)模型和數(shù)據(jù)集的規(guī)模增長速度已經(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ù)的語言模型（如 OpenAI 的大型生成預(yù)訓(xùn)練 Transformer 或最近類似的 BERT 模型）還是饋入 3000 萬個(gè)元素輸入的元學(xué)習(xí)神經(jīng)網(wǎng)絡(luò)（如我們在一篇 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ī)器的最簡單訓(xùn)練方法。

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

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

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

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

你要怎么做呢？

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

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

與之對等的 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ì)對 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

擴(kuò)展到極致

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

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