PyTorch -- 優化神經網絡訓練的 17 種方法

選自 efficientdl.com，作者：LORENZ KUHN

機器之心編譯

=============================================================================================================================================================================================

本文介紹在使用 PyTorch 訓練深度模型時最省力、最有效的 17 種方法。該文所提方法，都是假設你在 GPU 環境下訓練模型。具體內容如下。

17 種加速 PyTorch 訓練的方法

1. 考慮換一種學習率 schedule

學習率 schedule 的選擇對模型的收斂速度和泛化能力有很大的影響。Leslie N. Smith 等人在論文《Cyclical Learning Rates for Training Neural Networks》、《Super-Convergence: Very Fast Training of Neural Networks Using Large Learning Rates 》中提出了週期性（Cyclical）學習率以及 1Cycle 學習率 schedule。之後，fast.ai 的 Jeremy Howard 和 Sylvain Gugger 對其進行了推廣。下圖是 1Cycle 學習率 schedule 的圖示：

Sylvain 寫到：1Cycle 包括兩個等長的步幅，一個步幅是從較低的學習率到較高的學習率，另一個是回到最低水平。最大值來自學習率查找器選取的值，較小的值可以低十倍。然後，這個週期的長度應該略小於總的 epochs 數，並且，在訓練的最後階段，我們應該允許學習率比最小值小几個數量級。

與傳統的學習率 schedule 相比，在最好的情況下，該 schedule 實現了巨大的加速（Smith 稱之爲超級收斂）。例如，使用 1Cycle 策略在 ImageNet 數據集上訓練 ResNet-56，訓練迭代次數減少爲原來的 1/10，但模型性能仍能比肩原論文中的水平。在常見的體系架構和優化器中，這種 schedule 似乎表現得很好。

Pytorch 已經實現了這兩種方法：「torch.optim.lr_scheduler.CyclicLR」和「torch.optim.lr_scheduler.OneCycleLR」。

參考文檔：https://pytorch.org/docs/stable/optim.html

2. 在 DataLoader 中使用多個 worker 和頁鎖定內存

當使用 torch.utils.data.DataLoader 時，設置 num_workers > 0，而不是默認值 0，同時設置 pin_memory=True，而不是默認值 False。

參考文檔：https://pytorch.org/docs/stable/data.html

來自 NVIDIA 的高級 CUDA 深度學習算法軟件工程師 Szymon Micacz 就曾使用四個 worker 和頁鎖定內存（pinned memory）在單個 epoch 中實現了 2 倍的加速。人們選擇 worker 數量的經驗法則是將其設置爲可用 GPU 數量的四倍，大於或小於這個數都會降低訓練速度。請注意，增加 num_workers 將增加 CPU 內存消耗。

3. 把 batch 調到最大

把 batch 調到最大是一個頗有爭議的觀點。一般來說，如果在 GPU 內存允許的範圍內將 batch 調到最大，你的訓練速度會更快。但是，你也必須調整其他超參數，比如學習率。一個比較好用的經驗是，batch 大小加倍時，學習率也要加倍。

OpenAI 的論文《An Empirical Model of Large-Batch Training》很好地論證了不同的 batch 大小需要多少步才能收斂。在《How to get 4x speedup and better generalization using the right batch size》一文中，作者 Daniel Huynh 使用不同的 batch 大小進行了一些實驗（也使用上面討論的 1Cycle 策略）。最終，他將 batch 大小由 64 增加到 512，實現了 4 倍的加速。

然而，使用大 batch 的不足是，這可能導致解決方案的泛化能力比使用小 batch 的差。

4. 使用自動混合精度（AMP）

PyTorch 1.6 版本包括對 PyTorch 的自動混合精度訓練的本地實現。這裏想說的是，與單精度 (FP32) 相比，某些運算在半精度 (FP16) 下運行更快，而不會損失準確率。AMP 會自動決定應該以哪種精度執行哪種運算。這樣既可以加快訓練速度，又可以減少內存佔用。

在最好的情況下，AMP 的使用情況如下：

import torch
# Creates once at the beginning of training
scaler = torch.cuda.amp.GradScaler()


for data, label in data_iter:
   optimizer.zero_grad()
   # Casts operations to mixed precision
   with torch.cuda.amp.autocast():
      loss = model(data)

   # Scales the loss, and calls backward()
   # to create scaled gradients
   scaler.scale(loss).backward()

   # Unscales gradients and calls
   # or skips optimizer.step()
   scaler.step(optimizer)

   # Updates the scale for next iteration
   scaler.update()

5. 考慮使用另一種優化器

AdamW 是由 fast.ai 推廣的一種具有權重衰減（而不是 L2 正則化）的 Adam，在 PyTorch 中以 torch.optim.AdamW 實現。AdamW 似乎在誤差和訓練時間上都一直優於 Adam。

Adam 和 AdamW 都能與上面提到的 1Cycle 策略很好地搭配。

目前，還有一些非本地優化器也引起了很大的關注，最突出的是 LARS 和 LAMB。NVIDA 的 APEX 實現了一些常見優化器的融合版本，比如 Adam。與 PyTorch 中的 Adam 實現相比，這種實現避免了與 GPU 內存之間的多次傳遞，速度提高了 5%。

6. cudNN 基準

如果你的模型架構保持不變、輸入大小保持不變，設置 torch.backends.cudnn.benchmark = True。

7. 小心 CPU 和 GPU 之間頻繁的數據傳輸

當頻繁地使用 tensor.cpu() 將張量從 GPU 轉到 CPU（或使用 tensor.cuda() 將張量從 CPU 轉到 GPU）時，代價是非常昂貴的。item() 和 .numpy() 也是一樣可以使用. detach() 代替。

如果你創建了一個新的張量，可以使用關鍵字參數 device=torch.device(cuda:0) 將其分配給 GPU。

如果你需要傳輸數據，可以使用. to(non_blocking=True)，只要在傳輸之後沒有同步點。

8. 使用梯度 / 激活 checkpointing

Checkpointing 的工作原理是用計算換內存，並不存儲整個計算圖的所有中間激活用於 backward pass，而是重新計算這些激活。我們可以將其應用於模型的任何部分。

具體來說，在 forward pass 中，function 會以 torch.no_grad() 方式運行，不存儲中間激活。相反的是， forward pass 中會保存輸入元組以及 function 參數。在 backward pass 中，輸入和 function 會被檢索，並再次在 function 上計算 forward pass。然後跟蹤中間激活，使用這些激活值計算梯度。

因此，雖然這可能會略微增加給定 batch 大小的運行時間，但會顯著減少內存佔用。這反過來又將允許進一步增加所使用的 batch 大小，從而提高 GPU 的利用率。

儘管 checkpointing 以 torch.utils.checkpoint 方式實現，但仍需要一些思考和努力來正確地實現。Priya Goyal 寫了一個很好的教程來介紹 checkpointing 關鍵方面。

Priya Goyal 教程地址：

https://github.com/prigoyal/pytorch_memonger/blob/master/tutorial/Checkpointing_for_PyTorch_models.ipynb

9. 使用梯度積累

增加 batch 大小的另一種方法是在調用 optimizer.step() 之前在多個. backward() 傳遞中累積梯度。

Hugging Face 的 Thomas Wolf 的文章《Training Neural Nets on Larger Batches: Practical Tips for 1-GPU, Multi-GPU & Distributed setups》介紹瞭如何使用梯度累積。梯度累積可以通過如下方式實現：

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 accumulate

這個方法主要是爲了規避 GPU 內存的限制而開發的。

10. 使用分佈式數據並行進行多 GPU 訓練

加速分佈式訓練可能有很多方法，但是簡單的方法是使用 torch.nn.DistributedDataParallel 而不是 torch.nn.DataParallel。這樣一來，每個 GPU 將由一個專用的 CPU 核心驅動，避免了 DataParallel 的 GIL 問題。

分佈式訓練文檔地址：https://pytorch.org/tutorials/beginner/dist_overview.html

11. 設置梯度爲 None 而不是 0

梯度設置爲. zero_grad(set_to_none=True) 而不是 .zero_grad()。這樣做可以讓內存分配器處理梯度，而不是將它們設置爲 0。正如文檔中所說，將梯度設置爲 None 會產生適度的加速，但不要期待奇蹟出現。注意，這樣做也有缺點，詳細信息請查看文檔。

文檔地址：https://pytorch.org/docs/stable/optim.html

12. 使用. as_tensor() 而不是. tensor()

torch.tensor() 總是會複製數據。如果你要轉換一個 numpy 數組，使用 torch.as_tensor() 或 torch.from_numpy() 來避免複製數據。

13. 必要時打開調試工具

PyTorch 提供了很多調試工具，例如 autograd.profiler、autograd.grad_check、autograd.anomaly_detection。請確保當你需要調試時再打開調試器，不需要時要及時關掉，因爲調試器會降低你的訓練速度。

14. 使用梯度裁剪

關於避免 RNN 中的梯度爆炸的問題，已經有一些實驗和理論證實，梯度裁剪（gradient = min(gradient, threshold)）可以加速收斂。HuggingFace 的 Transformer 實現就是一個非常清晰的例子，說明了如何使用梯度裁剪。本文中提到的其他一些方法，如 AMP 也可以用。

在 PyTorch 中可以使用 torch.nn.utils.clip_grad_norm_來實現。

15. 在 BatchNorm 之前關閉 bias

在開始 BatchNormalization 層之前關閉 bias 層。對於一個 2-D 卷積層，可以將 bias 關鍵字設置爲 False：torch.nn.Conv2d(..., bias=False, ...)。

16. 在驗證期間關閉梯度計算

在驗證期間關閉梯度計算，設置：torch.no_grad() 。

17. 使用輸入和 batch 歸一化

要再三檢查一下輸入是否歸一化？是否使用了 batch 歸一化？

原文鏈接：https://efficientdl.com/faster-deep-learning-in-pytorch-a-guide/

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/fMyFmdIl7ikSbWXv2pT9tA