本文經AI新媒體量子位（公眾號ID:QbitAI）授權轉載，轉載請聯系出處。

如何提升PyTorch“煉丹”速度？

最近，有一位名叫Lorenz Kuhn的小哥，分享了他在煉丹過程中總結的17種投入最低、效果最好的提升訓練速度的方法，而且基本上都可以直接在PyTorch中進行更改，無需引入額外的庫。

不過需要注意的是，這些方法都是假設是在GPU上訓練模型。

這一分享在Reddit上得到了600的熱度。

接下來，我們便從提速高低開始，依次對這些方法來做介紹。

1、選擇合適的學習率時間表。

選擇的學習率時間表對收斂速度以及模型的泛化性能有很大影響。

Leslie Smith提出的周期性學習速率（CLR)以及 1cycle 策略可以令復雜模型的訓練迅速完成。

比如在 cifar10 上訓練 resnet-56 時，通過使用 1cycle，迭代次數可以減少10倍，得到與原論文相同的精度。

在最好的情況下，與傳統的相比，這個時間表實現了大規模的提速。不過有一個缺點，它們引入了一些額外的超參數。

為什么這樣做有效？一種可能的解釋是，定期增加學習率有助于更快地穿越損失函數中的鞍點。

2、在DataLoader中使用多個工作程序并固定內存。

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

英偉達高級工程師Szymon Micacz使用了4個工作程序和固定內存，在單個訓練時期內將速度提高了兩倍。

需要注意的是，在選擇worker數量時，建議將設置為可用GPU數量的四倍。

worker數量的多和少都會導致速度變慢，數量越多還會增加CPU內存消耗

3、批量最大化。

這一方法極具爭議。但在通常情況下，使用GPU內存允許的最大批處理量可以加快訓練速度。

如果要修改批量大小，還需要調整其他的超參數，比如，學習率。一般來說，將批量大小增加一倍，學習率也提高一倍。

此前有人進行了了一些不同批量大小的實驗，通過將批量大小從64增加到512實現了4倍的加速。

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

PyTorch 1.6版本就包括了對 PyTorch 的自動混合精度訓練的本地實現。

與其他地方使用的單精度（FP32）相比，某些操作可以在半精度（FP16）上運行得更快，并且不會損失準確性。

隨后，讓AMP自動決定應以什么樣的格式執行操作，這樣既可以加快訓練速度，也可以減少內存占用。

有研究者發現，在NVIDIA V100 GPU上對一些常見的語言和視覺模型進行基準測試時，使用AMP要比常規的FP32訓練的速度提升2倍，最高可提升5.5倍。

目前，只有CUDA ops 可以通過這種方式進行自動廣播。

5、使用不同的優化器

比如AdamW，AdamW是帶有權重衰減（而不是L2正則化）的Adam，它在錯誤實現、訓練時間都勝過Adam。

此外，還有一些非本地的優化器值得關注，比如，LARS和LAMB。

NVIDA的APEX實現了一些常見優化器（比如Adam）的融合版本，比如Adam。與Adam的PyTorch實現相比，它避免了多次進出GPU內存的過程，產生了5%左右的速度提升。

6、打開cudNN基準測試。

如果你的模型架構保持固定，輸入大小保持不變，則可以設置torch.backends.cudnn.benchmark = True，啟動 cudNN 自動調整器。

它將對cudNN中計算卷積的多種不同方法進行基準測試，以獲得最佳的性能指標。

7、防止CPU和GPU之間頻繁傳輸數據。

注意要經常使用tensor.cpu()將tensors從GPU傳輸到CPU，.item()和.numpy()也是如此，使用.detach()代替。

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

如果到傳輸數據的情境下，可以使用.to(non_blocking=True)，只要你在傳輸后沒有任何同步點。

8、使用梯度/激活檢查點。

檢查點的工作原理，是用計算換取內存。檢查點部分不是講整個計算圖的所有中間激活都存儲起來向后計算，而不是保存中間激活，在后傳中重新計算。

它可以應用到模型的任何部分。

具體來說，在前向傳遞中，函數將以torch.no_grad()的方式運行，即不存儲中間的激活。相反，前向傳遞會保存輸入元組和函數參數。

在后向傳遞中，檢索保存的輸入和函數，然后再次對函數進行前向傳遞計算，現在跟蹤中間激活，使用這些激活值計算梯度。

雖然這可能會略微增加你在給定批量大小下的運行時間，但你會顯著減少你的內存占用。這反過來又會讓你進一步增加你所使用的批次大小，提高GPU的利用率。

9、使用梯度累積。

另一種增加批次大小的方法是在調用optimizer.step()之前，在多個.backward()通道中累積梯度。

這個方法主要是為了規避GPU內存限制而開發的，但不清楚是否有額外的.backward()循環之間的權衡。

10、使用DistributedDataParallel進行多GPU訓練。

加速分布式訓練的方法可能需要單獨寫一篇文章，但一個簡單的方法是使用 torch.nn.DistributedDataParallel 而不是 torch.nn.DataParallel。

這樣做可以讓每個GPU將由一個專門的CPU核驅動，避免了DataParallel的GIL問題。

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、如果不需要，請關閉調試API。

Pytorch提供了很多調試工具，例如autograd.profiler，autograd.grad_check和autograd.anomaly_detection，確保在需要的時候使用它們，不需要時將其關閉，否則他們會拖慢你的訓練速度。

14、使用梯度剪裁。

剪裁梯度，可以加速加速收斂。最初是用來避免RNNs中的梯度爆炸，可以使用orch.nn.utils.clipgrad_norm來實現。

目前尚不清楚哪些模型能靠梯度剪裁能夠加速多少，但它似乎對RNNs、基于 Transformer 和 ResNets 的架構以及一系列不同的優化器都非常有用。

15、在BatchNorm之前關閉偏置。

這是一個非常簡單的方法，在BatchNormalization圖層之前關閉圖層的偏置。

對于二維卷積層，可以通過將bias關鍵字設置為False：來完成torch.nn.Conv2d(…, bias=False, …)

16、在驗證過程中關閉梯度計算。

在驗證期間設置torch.no_grad() 。

17、使用輸入和批次歸一化。

額外提示，使用JIT來融合逐點操作。

如果你有相鄰的逐點操作，可以使用PyTorch JIT將其合并成一個FusionGroup，然后在單個內核上啟動，這樣可以節省一些內存讀寫。

不少網友在表達感謝的同時，還分享了自己訓練時的小Tips。

比如這位煉丹師分享了第“18”個方法，下載更多的RAM。

還有人提出了兩點建議：

1、數據變換 (用于數據增強) 可成為速度提升的另一個來源。一些只使用簡單 Python 語句的變換可以通過使用 numba 包來加速。

2、將數據集預處理成單個文件，對速度也有好處。

除了這些，你還有哪些可以提升訓練速度的方法？歡迎與我們分享~