在深度學習訓練中，我們經常遇到 GPU 的內存太小的問題，如果我們的數據量比較大，別說大批量（large batch size）訓練了，有時候甚至連一個訓練樣本都放不下。但是隨機梯度下降（SGD）中，如果能使用更大的 Batch Size 訓練，一般能得到更好的結果。所以問題來了：
問題來了：當 GPU 的內存不夠時，如何使用大批量（large batch size）樣本來訓練神經網絡呢？
這篇文章將以 PyTorch 為例，講解一下幾點：
- 當 GPU 的內存小於 Batch Size 的訓練樣本，或者甚至連一個樣本都塞不下的時候，怎么用單個或多個 GPU 進行訓練？
- 怎么盡量高效地利用多 GPU？
單個或多個 GPU 進行大批量訓練
如果你也遇到過 CUDA RuntimeError: out of memory 的錯誤，那么說明你也遇到了這個問題。
PyTorch 的開發人員都出來了，估計一臉黑線：兄弟，這不是 bug，是你內存不夠…
又一個方法可以解決這個問題：梯度累加（accumulating gradients）。
一般在 PyTorch 中，我們是這樣來更新梯度的：
|
1
2
3
4
5
|
predictions = model(inputs) # 前向計算
loss = loss_function(predictions, labels) # 計算損失函數
loss.backward() # 后向計算梯度
optimizer.step() # 優化器更新梯度
predictions = model(inputs) # 用更新過的參數值進行下一次前向計算
|
在上看的代碼注釋中，在計算梯度的 loss.backward() 操作中，每個參數的梯度被計算出來后，都被存儲在各個參數對應的一個張量里：parameter.grad。然后優化器就會根據這個來更新每個參數的值，就是 optimizer.step()。
而梯度累加（accumulating gradients）的基本思想就是， 在優化器更新參數前，也就是執行 optimizer.step() 前，我們進行多次梯度計算，保存在 parameter.grad 中，然后累加梯度再更新。這個在 PyTorch 中特別容易實現，因為 PyTorch 中，梯度值本身會保留，除非我們調用 model.zero_grad() or optimizer.zero_grad()。
下面是一個梯度累加的例子，其中 accumulation_steps 就是要累加梯度的循環數：
|
1
2
3
4
5
6
7
8
9
|
model.zero_grad() # 重置保存梯度值的張量
for i, (inputs, labels) in enumerate(training_set):
predictions = model(inputs) # 前向計算
loss = loss_function(predictions, labels) # 計算損失函數
loss = loss / accumulation_steps # 對損失正則化 (如果需要平均所有損失)
loss.backward() # 計算梯度
if (i 1) % accumulation_steps == 0: # 重復多次前面的過程
optimizer.step() # 更新梯度
model.zero_grad() # 重置梯度
|
如果連一個樣本都不放下怎么辦？
如果樣本特別大，別說 batch training，要是 GPU 的內存連一個樣本都不下怎么辦呢？
答案是使用梯度檢查點（gradient-checkpoingting），用計算量來換內存。基本思想就是，在反向傳播的過程中，把梯度切分成幾部分，分別對網絡上的部分參數進行更新（見下圖）。但這種方法的速度很慢，因為要增加額外的計算量。但在某些例子上又很有用，比如訓練長序列的 RNN 模型等（感興趣的話可以參考這篇文章）。
圖片來自：https://medium.com/tensorflow/fitting-larger-networks-into-memory-583e3c758ff9
這里就不展開講了，可以參考 PyTorch 官方文檔對 Checkpoint 的描述：https://pytorch.org/docs/stable/checkpoint.html
多 GPU 訓練方法
簡單來講，PyTorch 中多 GPU 訓練的方法是使用 torch.nn.DataParallel。非常簡單，只需要一行代碼：
|
1
2
3
4
5
6
7
|
parallel_model = torch.nn.DataParallel(model) # 就是這里！
predictions = parallel_model(inputs) # 前向計算
loss = loss_function(predictions, labels) # 計算損失函數
loss.mean().backward() # 計算多個GPU的損失函數平均值，計算梯度
optimizer.step() # 反向傳播
predictions = parallel_model(inputs)
|
在使用torch.nn.DataParallel 的過程中，我們經常遇到一個問題：第一個GPU的計算量往往比較大。我們先來看一下多 GPU 的訓練過程原理：
在上圖第一行第四個步驟中，GPU-1 其實匯集了所有 GPU 的運算結果。這個對於多分類問題還好，但如果是自然語言處理模型就會出現問題，導致 GPU-1 匯集的梯度過大，直接爆掉。
那么就要想辦法實現多 GPU 的負載均衡，方法就是讓 GPU-1 不匯集梯度，而是保存在各個 GPU 上。這個方法的關鍵就是要分布化我們的損失函數，讓梯度在各個 GPU 上單獨計算和反向傳播。這里又一個開源的實現：https://github.com/zhanghang1989/PyTorch-Encoding。這里是一個修改版，可以直接在我們的代碼里調用：地址。實例：
|
1
2
3
4
5
6
7
8
9
10
11
|
from parallel import DataParallelModel, DataParallelCriterion
parallel_model = DataParallelModel(model) # 並行化model
parallel_loss = DataParallelCriterion(loss_function) # 並行化損失函數
predictions = parallel_model(inputs) # 並行前向計算
# "predictions"是多個gpu的結果的元組
loss = parallel_loss(predictions, labels) # 並行計算損失函數
loss.backward() # 計算梯度
optimizer.step() # 反向傳播
predictions = parallel_model(inputs)
|
如果你的網絡輸出是多個，可以這樣分解：
|
1
|
output_1, output_2 = zip(*predictions)
|
如果有時候不想進行分布式損失函數計算，可以這樣手動匯集所有結果：
|
1
|
gathered_predictions = parallel.gather(predictions)
|
下圖展示了負載均衡以后的原理：
（原文鏈接： https://medium.com/huggingface/training-larger-batches-practical-tips-on-1-gpu-multi-gpu-distributed-setups-ec88c3e51255）