如何提高PyTorch“煉丹”速度？這位小哥總結(jié)了17種方法，可直接上手更改的那種

白交 2021-01-15 10:18:05 來源：量子位

楊凈 發(fā)自 凹非寺
量子位 報道 | 公眾號 QbitAI

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

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

不過需要注意的是，這些方法都是假設(shè)是在GPU上訓(xùn)練模型。

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

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

1、選擇合適的學(xué)習(xí)率時間表。

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

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

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

在最好的情況下，與傳統(tǒng)的相比，這個時間表實現(xiàn)了大規(guī)模的提速。不過有一個缺點，它們引入了一些額外的超參數(shù)。

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

2、在DataLoader中使用多個工作程序并固定內(nèi)存。

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

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

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

worker數(shù)量的多和少都會導(dǎo)致速度變慢，數(shù)量越多還會增加CPU內(nèi)存消耗

3、批量最大化。

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

如果要修改批量大小，還需要調(diào)整其他的超參數(shù)，比如，學(xué)習(xí)率。一般來說，將批量大小增加一倍，學(xué)習(xí)率也提高一倍。

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

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

PyTorch 1.6版本就包括了對 PyTorch 的自動混合精度訓(xùn)練的本地實現(xiàn)。

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

隨后，讓AMP自動決定應(yīng)以什么樣的格式執(zhí)行操作，這樣既可以加快訓(xùn)練速度，也可以減少內(nèi)存占用。

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

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

5、使用不同的優(yōu)化器

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

此外，還有一些非本地的優(yōu)化器值得關(guān)注，比如，LARS和LAMB。

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

6、打開cudNN基準(zhǔn)測試。

如果你的模型架構(gòu)保持固定，輸入大小保持不變，則可以設(shè)置torch.backends.cudnn.benchmark = True，啟動 cudNN 自動調(diào)整器。

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

7、防止CPU和GPU之間頻繁傳輸數(shù)據(jù)。

注意要經(jīng)常使用tensor.cpu()將tensors從GPU傳輸?shù)紺PU，.item()和.numpy()也是如此，使用.detach()代替。

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

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

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

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

它可以應(yīng)用到模型的任何部分。

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

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

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

9、使用梯度累積。

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

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

10、使用DistributedDataParallel進行多GPU訓(xùn)練。

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

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

11、將梯度設(shè)置為None而不是0。

使用.zero_grad(set_to_none=True)而不是.zero_grad()。

這樣做會讓內(nèi)存分配器來處理梯度，而不是主動將它們設(shè)置為0，這樣會適度加速。

注意，這樣做并不是沒有副作用的。

12、使用 .as_tensor 而不是 .tensor()

torch.tensor() 總是復(fù)制數(shù)據(jù)。如果你有一個要轉(zhuǎn)換的 numpy 數(shù)組，使用 torch.as_tensor() 或 torch.from_numpy() 來避免復(fù)制數(shù)據(jù)。

13、如果不需要，請關(guān)閉調(diào)試API。

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

14、使用梯度剪裁。

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

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

15、在BatchNorm之前關(guān)閉偏置。

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

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

16、在驗證過程中關(guān)閉梯度計算。

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

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

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

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

不少網(wǎng)友在表達感謝的同時，還分享了自己訓(xùn)練時的小Tips。

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

還有人提出了兩點建議：

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

2、將數(shù)據(jù)集預(yù)處理成單個文件，對速度也有好處。

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

傳送門：
https://efficientdl.com/faster-deep-learning-in-pytorch-a-guide/#1-consider-using-another-learning-rate-schedule
https://www.reddit.com/r/MachineLearning/comments/kvs1ex/d_here_are_17_ways_of_making_pytorch_training/