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

白交 2021-01-15 10:18:05 來(lái)源：量子位

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

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

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

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

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

接下來(lái)，我們便從提速高低開(kāi)始，依次對(duì)這些方法來(lái)做介紹。

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

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

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

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

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

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

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

使用時(shí)torch.utils.data.DataLoader，請(qǐng)?jiān)O(shè)置num_workers > 0，而不是默認(rèn)值0，和pin_memory=True，而不是默認(rèn)值False。

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

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

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

3、批量最大化。

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

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

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

4、使用自動(dòng)混合精度（AMP）。

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

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

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

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

目前，只有CUDA ops 可以通過(guò)這種方式進(jìn)行自動(dòng)廣播。

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

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

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

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

6、打開(kāi)cudNN基準(zhǔn)測(cè)試。

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

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

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

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

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

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

8、使用梯度/激活檢查點(diǎn)。

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

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

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

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

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

9、使用梯度累積。

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

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

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

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

這樣做可以讓每個(gè)GPU將由一個(gè)專門(mén)的CPU核驅(qū)動(dòng)，避免了DataParallel的GIL問(wèn)題。

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

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

這樣做會(huì)讓內(nèi)存分配器來(lái)處理梯度，而不是主動(dòng)將它們?cè)O(shè)置為0，這樣會(huì)適度加速。

注意，這樣做并不是沒(méi)有副作用的。

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

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

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

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

14、使用梯度剪裁。

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

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

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

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

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

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

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

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

額外提示，使用JIT來(lái)融合逐點(diǎn)操作。

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

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

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

還有人提出了兩點(diǎn)建議：

1、數(shù)據(jù)變換 (用于數(shù)據(jù)增強(qiáng)) 可成為速度提升的另一個(gè)來(lái)源。一些只使用簡(jiǎn)單 Python 語(yǔ)句的變換可以通過(guò)使用 numba 包來(lái)加速。

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

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

傳送門(mé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/