如何提高PyTorch“煉丹”速度？這位小哥總結(jié)了17種方法，可直接上手更改的那種
本文經(jīng)AI新媒體量子位（公眾號ID:QbitAI）授權(quán)轉(zhuǎn)載，轉(zhuǎn)載請聯(lián)系出處。
如何提升PyTorch“煉丹”速度？
最近，有一位名叫Lorenz Kuhn的小哥，分享了他在煉丹過程中總結(jié)的17種投入最低、效果最好的提升訓(xùn)練速度的方法，而且基本上都可以直接在PyTorch中進(jìn)行更改，無需引入額外的庫。
不過需要注意的是，這些方法都是假設(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)的相比，這個時間表實(shí)現(xiàn)了大規(guī)模的提速。不過有一個缺點(diǎn)，它們引入了一些額外的超參數(shù)。
為什么這樣做有效？一種可能的解釋是，定期增加學(xué)習(xí)率有助于更快地穿越損失函數(shù)中的鞍點(diǎn)。
2、在DataLoader中使用多個工作程序并固定內(nèi)存。
使用時torch.utils.data.DataLoader，請?jiān)O(shè)置num_workers > 0，而不是默認(rèn)值0，和pin_memory=True，而不是默認(rèn)值False。
英偉達(dá)高級工程師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í)率也提高一倍。
此前有人進(jìn)行了了一些不同批量大小的實(shí)驗(yàn)，通過將批量大小從64增加到512實(shí)現(xiàn)了4倍的加速。
4、使用自動混合精度（AMP）。
PyTorch 1.6版本就包括了對 PyTorch 的自動混合精度訓(xùn)練的本地實(shí)現(xiàn)。
與其他地方使用的單精度（FP32）相比，某些操作可以在半精度（FP16）上運(yùn)行得更快，并且不會損失準(zhǔn)確性。
隨后，讓AMP自動決定應(yīng)以什么樣的格式執(zhí)行操作，這樣既可以加快訓(xùn)練速度，也可以減少內(nèi)存占用。
有研究者發(fā)現(xiàn)，在NVIDIA V100 GPU上對一些常見的語言和視覺模型進(jìn)行基準(zhǔn)測試時，使用AMP要比常規(guī)的FP32訓(xùn)練的速度提升2倍，最高可提升5.5倍。
目前，只有CUDA ops 可以通過這種方式進(jìn)行自動廣播。
5、使用不同的優(yōu)化器
比如AdamW，AdamW是帶有權(quán)重衰減（而不是L2正則化）的Adam，它在錯誤實(shí)現(xiàn)、訓(xùn)練時間都勝過Adam。
此外，還有一些非本地的優(yōu)化器值得關(guān)注，比如，LARS和LAMB。
NVIDA的APEX實(shí)現(xiàn)了一些常見優(yōu)化器（比如Adam）的融合版本，比如Adam。與Adam的PyTorch實(shí)現(xiàn)相比，它避免了多次進(jìn)出GPU內(nèi)存的過程，產(chǎn)生了5%左右的速度提升。
6、打開cudNN基準(zhǔn)測試。
如果你的模型架構(gòu)保持固定，輸入大小保持不變，則可以設(shè)置torch.backends.cudnn.benchmark = True，啟動 cudNN 自動調(diào)整器。
它將對cudNN中計(jì)算卷積的多種不同方法進(jìn)行基準(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)，只要你在傳輸后沒有任何同步點(diǎn)。
8、使用梯度/激活檢查點(diǎn)。
檢查點(diǎn)的工作原理，是用計(jì)算換取內(nèi)存。檢查點(diǎn)部分不是講整個計(jì)算圖的所有中間激活都存儲起來向后計(jì)算，而不是保存中間激活，在后傳中重新計(jì)算。
它可以應(yīng)用到模型的任何部分。
具體來說，在前向傳遞中，函數(shù)將以torch.no_grad()的方式運(yùn)行，即不存儲中間的激活。相反，前向傳遞會保存輸入元組和函數(shù)參數(shù)。
在后向傳遞中，檢索保存的輸入和函數(shù)，然后再次對函數(shù)進(jìn)行前向傳遞計(jì)算，現(xiàn)在跟蹤中間激活，使用這些激活值計(jì)算梯度。
雖然這可能會略微增加你在給定批量大小下的運(yùn)行時間，但你會顯著減少你的內(nèi)存占用。這反過來又會讓你進(jìn)一步增加你所使用的批次大小，提高GPU的利用率。
9、使用梯度累積。
另一種增加批次大小的方法是在調(diào)用optimizer.step()之前，在多個.backward()通道中累積梯度。
這個方法主要是為了規(guī)避GPU內(nèi)存限制而開發(fā)的，但不清楚是否有額外的.backward()循環(huán)之間的權(quán)衡。
10、使用DistributedDataParallel進(jìn)行多GPU訓(xùn)練。
加速分布式訓(xùn)練的方法可能需要單獨(dú)寫一篇文章，但一個簡單的方法是使用 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來實(shí)現(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、在驗(yàn)證過程中關(guān)閉梯度計(jì)算。
在驗(yàn)證期間設(shè)置torch.no_grad() 。
17、使用輸入和批次歸一化。
額外提示，使用JIT來融合逐點(diǎn)操作。
如果你有相鄰的逐點(diǎn)操作，可以使用PyTorch JIT將其合并成一個FusionGroup，然后在單個內(nèi)核上啟動，這樣可以節(jié)省一些內(nèi)存讀寫。
不少網(wǎng)友在表達(dá)感謝的同時，還分享了自己訓(xùn)練時的小Tips。
比如這位煉丹師分享了第“18”個方法，下載更多的RAM。
還有人提出了兩點(diǎn)建議：
1、數(shù)據(jù)變換 (用于數(shù)據(jù)增強(qiáng)) 可成為速度提升的另一個來源。一些只使用簡單 Python 語句的變換可以通過使用 numba 包來加速。
2、將數(shù)據(jù)集預(yù)處理成單個文件，對速度也有好處。
除了這些，你還有哪些可以提升訓(xùn)練速度的方法？歡迎與我們分享~
