PyTorch 指南：17個技巧讓你的深度學習模型訓練變得飛快！
如果你正在 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策略。
此外，LARS和LAMB等其他優化器也收到廣泛關注。
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倍。
本文轉自雷鋒網，如需轉載請至雷鋒網官網申請授權。
