PyTorch 指南：17個技巧讓你的深度學習模型訓練變得飛快！
- 2021 年 1 月 26 日
- AI
雙語原文鏈接：Faster Deep Learning Training with PyTorch – a 2021 Guide
如果你正在 pytorch 中訓練深度學習模型，那麼如何能夠加快模型訓練速度呢？
在本文中，我會介紹一些改動最小、影響最大的在pytorch中加速深度學習模型的方法。對於每種方法，我會對其思路進行簡要介紹，然後預估提升速度並討論其限制。我會把我認為重要的部分強調介紹，並在每個部分展示一些實例。接下來我將假設你正在使用GPU訓練模型，這些方法基本不需要導入其他的庫，只需要再pytorch內進行更改即可。
以下是我根據預估的加速效果對不同方法的排序：
-
考慮使用其他的學習率調整計劃
-
在DataLoader中使用多個輔助進程並頁鎖定記憶體
-
最大化batch大小
-
使用自動混合精度AMP
-
考慮不同的優化器
-
打開cudNN基準
-
當心CPU與GPU之間的數據傳輸
-
使用梯度/激活檢查點
-
使用梯度累積
-
多GPU分散式訓練
-
將梯度設置為None而不是0
-
使用.as_tensor()而不是.tensor()
-
只在需要的時候打開debugging模式
-
使用梯度裁剪
-
在BatchNorm之前忽略偏差
-
驗證時關閉梯度計算
-
規範化輸入和批處理
1. 考慮使用其他的學習率調整計劃
在訓練中使用的學習率調整計劃會極大影響收斂速率以及模型泛化能力。
Leslie N. Smith 提出了循環學習率和1Cycle 學習率方法，然後由 fast.ai 的 Jeremy Howard 和 Sylvain Gugger 推廣了。總的來說，1Cycle 學習速率方法如下圖所示：
在最好的情況下，與傳統的學習率策略相比，這種策略可以實現巨大的加速—— Smith稱之為「超級收斂」。例如，使用1Cycle策略，在ImageNet上減少了ResNet-56訓練迭代數的10倍，就可以匹配原始論文的性能。該策略似乎在通用架構和優化器之間運行得很好。
PyTorch提供了 torch.optim.lr_scheduler.CyclicLR 和 torch.optim.lr_scheduler.OneCycleLR 兩種方法實現該操作，請參閱相關文檔。
這兩個方法的一個缺點是引入了許多額外的超參數。這篇文章和倉庫對如何查找好的超參數（包括上文提及的學習率）提供了詳細概述和實現。
至於為什麼要這樣做？現今並不完全清楚，但一個可能的解釋是：定期提高學習率有助於更快越過損失鞍點。
2. 在DataLoader中使用多個輔助進程並頁鎖定記憶體
在使用 torch.utils.data.DataLoader時，令 num_workers > 0，而不是默認值 0，同時設置 pin_memory=True，而不是默認值 False。至於為什麼這麼做，這篇文章會給你答案。
根據上述方法，Szymon Micacz 在四個 worker 和頁鎖定記憶體的情況下，在單個epoch中實現了 2 倍加速。
根據經驗，一般將進程數量設置為可用 GPU 數量的四倍，大於或小於這個值都會降低訓練速度。但是要注意，增加num_workers會增加 CPU 記憶體消耗。
3.最大化batch大小
一直以來，人們對於調大batch沒有定論。一般來說，在GPU記憶體允許的情況下增大batch將會增快訓練速度，但同時還需要調整學習率等其他超參數。根據經驗，batch大小加倍時，學習率也相應加倍。
OpenAI 的論文表明不同的batch大小收斂周期不同。Daniel Huynh用不同的batch大小進行了一些實驗（使用上述1Cycle 策略），實驗中他將 batch大小由64增加到512，實現了4倍加速。
然而也要注意，較大的batch會降低模型泛化能力，反之亦然。
4. 使用自動混合精度AMP
PyTorch1.6支援本地自動混合精度訓練。與單精度 (FP32) 相比，一些運算在不損失準確率的情況下，使用半精度 (FP16)速度更快。AMP能夠自動決定應該以哪種精度執行哪種運算，這樣既可以加快訓練速度，又減少了記憶體佔用。
AMP的使用如下所示：
| import torch# Creates once at the beginning of trainingscaler = 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() |
Huang及其同事在NVIDIA V100 GPU上對一些常用語言和視覺模型進行了基準測試，發現在FP32訓練中使用AMP提高約2倍的訓練速度，最高甚至達到5.5倍。
目前，只有CUDA支援上述方式，查看本文檔了解更多資訊。
5. 考慮不同的優化器
AdamW是由fast.ai提出的具有權重衰減（而非 L2 正則化）的Adam， PyTorch中通過torch.optim.AdamW實現。在誤差和訓練時間上，AdamW都優於Adam。查看此文章了解為什麼權重衰減使得Adam產生更好效果。
Adam和AdamW都很適合前文提到的1Cycle策略。
NVIDA的APEX對Adam等常見優化器進行優化融合，相比PyTorch中的原始Adam，由於避免了GPU記憶體之間的多次傳遞，訓練速度提升約 5%。
6. 打開cudNN基準
如果你的模型架構時固定的，同時輸入大小保持不變，那麼設置torch.backends.cudnn.benchmark = True可能會提升模型速度（幫助文檔）。通過啟用cudNN自動調節器，可以在cudNN中對多種計算卷積的方法進行基準測試，然後選擇最快的方法。
至於提速效果，Szymon Migacz在前向卷積時提速70％，在同時向前和後向卷積時提升了27％。
注意，如果你想要根據上述方法最大化批大小，該自動調整可能會非常耗時。
7. 當心CPU與GPU之間的數據傳輸
通過tensor.cpu()可以將張量從GPU傳輸到CPU，反之使用tensor.cuda()，但這樣的數據轉化代價較高。 .item()和.numpy()的使用也是如此，建議使用.detach（）。
如果要創建新的張量，使用關鍵字參數device=torch.device（’cuda：0’）將其直接分配給GPU。
最好使用.to(non_blocking=True)傳輸數據，確保傳輸後沒有任何同步點即可。
另外Santosh Gupta的SpeedTorch也值得一試，儘管其加速與否尚不完全清除。
8.使用梯度/激活檢查點
檢查點通過將計算保存到記憶體來工作。檢查點在反向傳播演算法過程中並不保存計算圖的中間激活，而是在反向傳播時重新計算，其可用於模型的任何部分。
具體來說，在前向傳播中，function以torch.no_grad()方式運行，不存儲任何中間激活。相反，前向傳遞將保存輸入元組和function參數。在反向傳播時，檢索保存的輸入和function，並再次對function進行正向傳播，記錄中間激活，並使用這些激活值計算梯度。
因此，對於特定的批處理大小，這可能會稍微增加運行時間，但會顯著減少記憶體消耗。反過來，你可以進一步增加批處理大小，從而更好地利用GPU。
雖然檢查點可以通過torch.utils.checkpoint方便實現，但仍需要里哦阿姐其思想與本質。Priya Goyal的教程很清晰的演示了檢查點的一些關鍵思想，推薦閱讀。
9.使用梯度累積
增加批處理大小的另一種方法是在調用Optimizer.step()之對多個.backward()傳遞梯度進行累積。
根據Hugging Face的Thomas Wolf發表的文章，可以按以下方式實現梯度累積：
| 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 |
該方法主要是為了規避GPU記憶體的限制，但對其他.backward()循環之間的取捨我並不清楚。fastai論壇上的討論似乎表明它實際上是可以加速訓練的，因此值得一試。詳情查看GitHub託管的rawgradient_accumulation.py。
10.多GPU分散式訓練
通過分散式訓練加快模型速度的一種簡單的方法是使用torch.nn.DistributedDataParallel而不是torch.nn.DataParallel。這樣，每個GPU將由專用的CPU內核驅動，從而避免了DataParallel的GIL問題。
強烈推薦閱讀分散式訓練相關文檔了解更多資訊：
PyTorch Distributed Overview — PyTorch Tutorials 1.7.0 documentation
11.將梯度設置為None而不是0
設置.zero_grad(set_to_none=True)而不是.zero_grad()。
這樣記憶體分配器處理梯度而不是主動將其設置為0，這會產生該文檔所示的適度加速，但不要抱有過大期望。
注意，這樣做不會有任何副作用！閱讀文檔查看更多資訊。
12.使用.as_tensor()而不是.tensor()
torch.tensor()本質是複製數據，因此，如果要轉換numpy數組，使用torch.as_tensor()或torch.from_numpy()可以避免複製數據。
13.只在需要的時候打開debugging模式
Pytorch提供了許多調試工具，例如autograd.profiler, autograd.grad_check和autograd.anomaly_detection。使用時一定要謹慎，這些調試工具顯然會影響訓練速度，因此在不需要時將其關閉。
14.使用梯度裁剪
為了避免RNN中的梯度爆炸，使用梯度裁剪gradient = min(gradient, threshold)可以起到加速收斂作用，這一方法已得到理論和實驗的支援。
Hugging Face的Transformer提供了將梯度裁剪和AMP等其他方法有效結合的清晰示例。
在PyTorch中，也可使用torch.nn.utils.clip_grad_norm_（文檔查閱）完成此操作。
雖然我尚不完全清楚哪種模型可以從梯度裁剪中受益，但毫無疑問的是，對於RNN、基於Transformer和ResNets結構的一系列優化器來說，該方法顯然是起到一定作用的。
15.在BatchNorm之前忽略偏差
在BatchNormalization層之前關閉之前層的偏差時一種簡單有效的方法。對於二維卷積層，可以通過將bias關鍵字設置為False實現，即torch.nn.Conv2d(…, bias=False, …)。閱讀該文檔了解其原理。
與其他方法相比，該方法的速度提升是有的。
16. 驗證時關閉梯度計算
在模型驗證時令torch.no_grad()
17. 規範化輸入和批處理
也許你已經在這樣做了，但還是要仔細檢查，反覆確認：
點擊查看這樣做的原因。
其他技巧：使用JIT實現逐點融合
如果要執行相鄰逐點操作，可以使用PyTorch JIT將它們組合成一個FusionGroup，然後在單內核上啟動，而不是像默認情況那樣在多個內核上啟動，同時還可以保存一些記憶體進行讀寫。
Szymon Migacz展示了如何使用@torch.jit.script裝飾器融合GELU操作融合，如下：
| @torch.jit.scriptdef fused_gelu(x): return x * 0.5 * (1.0 + torch.erf(x / 1.41421)) |
相比於未融合版本，融合這些操作可以使fused_gelu的執行速度提高5倍。
查閱此文章獲取更多使用Torchscript加速RNN的示例。
當然，你還可以在Reddit上與u/Patient_Atmosphere45交流討論。
參考及其他資源
本文許多技巧參考自Szymon Migacz的演講及PyTorch文檔。
PyTorch Lightning的作者William Falcon在這兩篇文章種介紹了關於加快訓練的內容。同時，PyTorch Lightning已集成以上一些技巧與方法。
Hugging Face的作者Thomas Wolf也寫了一系列文章介紹深度學習的加速-尤其是語言模型。
Sylvain Gugger和Jeremy Howard寫了很多關於學習率和AdamW的文章。
感謝Ben Hahn，Kevin Klein和Robin Vaaler對本文撰寫提供的幫助！
AI研習社是AI學術青年和AI開發者技術交流的在線社區。我們與高校、學術機構和產業界合作，通過提供學習、實戰和求職服務，為AI學術青年和開發者的交流互助和職業發展打造一站式平台，致力成為中國最大的科技創新人才聚集地。
如果，你也是位熱愛分享的AI愛好者。歡迎與譯站一起，學習新知，分享成長。
