來源：小白學視覺

掌握這 17 種方法，用最省力的方式，加速你的 Pytorch 深度學習訓練。

近日，Reddit 上一個帖子熱度爆表。主題內容是關于怎樣加速 PyTorch 訓練。原文作者是來自蘇黎世聯邦理工學院的計算機科學碩士生 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」。

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 關鍵方面。

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 accumulat

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

10. 使用分布式數據并行進行多 GPU 訓練

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

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

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

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 歸一化？