近日，Reddit 上一個帖子熱度爆表。主題內容是關於怎樣加速 PyTorch 訓練。原文作者是來自蘇黎世聯邦理工學院的電腦科學碩士生 LORENZ KUHN，文章向我們介紹了在使用 PyTorch 訓練深度模型時最省力、最有效的 17 種方法。
該文所提方法，都是假設你在 GPU 環境下訓練模型。具體內容如下。學習率 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.html2. 在 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 記憶體消耗。把 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 的差。PyTorch 1.6 版本包括對 PyTorch 的自動混合精度訓練的本地實現。這裡想說的是，與單精度 (FP32) 相比，某些運算在半精度 (FP16) 下執行更快，而不會損失準確率。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()
AdamW 是由 fast.ai 推廣的一種具有權重衰減（而不是 L2 正則化）的 Adam，在 PyTorch 中以 torch.optim.AdamW 實現。AdamW 似乎在誤差和訓練時間上都一直優於 Adam。Adam 和 AdamW 都能與上面提到的 1Cycle 策略很好地搭配。目前，還有一些非本地最佳化器也引起了很大的關注，最突出的是 LARS 和 LAMB。NVIDA 的 APEX 實現了一些常見最佳化器的融合版本，比如 Adam。與 PyTorch 中的 Adam 實現相比，這種實現避免了與 GPU 記憶體之間的多次傳遞，速度提高了 5%。如果你的模型架構保持不變、輸入大小保持不變，設定 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. 使用梯度 / 啟用 checkpointingCheckpointing 的工作原理是用計算換記憶體，並不儲存整個計算圖的所有中間啟用用於 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 關鍵方面。https://github.com/prigoyal/pytorch_memonger/blob/master/tutorial/Checkpointing_for_PyTorch_models.ipynb增加 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 tensorsfor 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 記憶體的限制而開發的。加速分散式訓練可能有很多方法，但是簡單的方法是使用 torch.nn.DistributedDataParallel 而不是 torch.nn.DataParallel。這樣一來，每個 GPU 將由一個專用的 CPU 核心驅動，避免了 DataParallel 的 GIL 問題。分散式訓練文件地址：https://pytorch.org/tutorials/beginner/dist_overview.html梯度設定為. zero_grad(set_to_none=True) 而不是 .zero_grad()。這樣做可以讓記憶體分配器處理梯度，而不是將它們設定為 0。正如文件中所說，將梯度設定為 None 會產生適度的加速，但不要期待奇蹟出現。注意，這樣做也有缺點，詳細資訊請檢視文件。文件地址：https://pytorch.org/docs/stable/optim.html12. 使用. as_tensor() 而不是. tensor()torch.tensor() 總是會複製資料。如果你要轉換一個 numpy 陣列，使用 torch.as_tensor() 或 torch.from_numpy() 來避免複製資料。PyTorch 提供了很多除錯工具，例如 autograd.profiler、autograd.grad_check、autograd.anomaly_detection。請確保當你需要除錯時再開啟偵錯程式，不需要時要及時關掉，因為偵錯程式會降低你的訓練速度。關於避免 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, ...)。在驗證期間關閉梯度計算，設定：torch.no_grad() 。要再三檢查一下輸入是否歸一化？是否使用了 batch 歸一化？原文連結：https://efficientdl.com/faster-deep-learning-in-pytorch-a-guide/