0
| 本文作者： AI研習(xí)社-譯站 | 2021-01-26 11:27 |
雙語原文鏈接：Faster Deep Learning Training with PyTorch – a 2021 Guide
如果你正在 pytorch 中訓(xùn)練深度學(xué)習(xí)模型，那么如何能夠加快模型訓(xùn)練速度呢？
在本文中，我會介紹一些改動最小、影響最大的在pytorch中加速深度學(xué)習(xí)模型的方法。對于每種方法，我會對其思路進(jìn)行簡要介紹，然后預(yù)估提升速度并討論其限制。我會把我認(rèn)為重要的部分強(qiáng)調(diào)介紹，并在每個部分展示一些實例。接下來我將假設(shè)你正在使用GPU訓(xùn)練模型，這些方法基本不需要導(dǎo)入其他的庫，只需要再pytorch內(nèi)進(jìn)行更改即可。
以下是我根據(jù)預(yù)估的加速效果對不同方法的排序：
考慮使用其他的學(xué)習(xí)率調(diào)整計劃
在DataLoader中使用多個輔助進(jìn)程并頁鎖定內(nèi)存
最大化batch大小
使用自動混合精度AMP
考慮不同的優(yōu)化器
打開cudNN基準(zhǔn)
當(dāng)心CPU與GPU之間的數(shù)據(jù)傳輸
使用梯度/激活檢查點
使用梯度累積
多GPU分布式訓(xùn)練
將梯度設(shè)置為None而不是0
使用.as_tensor()而不是.tensor()
只在需要的時候打開debugging模式
使用梯度裁剪
在BatchNorm之前忽略偏差
驗證時關(guān)閉梯度計算
規(guī)范化輸入和批處理
在訓(xùn)練中使用的學(xué)習(xí)率調(diào)整計劃會極大影響收斂速率以及模型泛化能力。
Leslie N. Smith 提出了循環(huán)學(xué)習(xí)率和1Cycle 學(xué)習(xí)率方法，然后由 fast.ai 的 Jeremy Howard 和 Sylvain Gugger 推廣了?？偟膩碚f，1Cycle 學(xué)習(xí)速率方法如下圖所示：
在最好的情況下，與傳統(tǒng)的學(xué)習(xí)率策略相比，這種策略可以實現(xiàn)巨大的加速—— Smith稱之為“超級收斂”。例如，使用1Cycle策略，在ImageNet上減少了ResNet-56訓(xùn)練迭代數(shù)的10倍，就可以匹配原始論文的性能。該策略似乎在通用架構(gòu)和優(yōu)化器之間運(yùn)行得很好。
PyTorch提供了 torch.optim.lr_scheduler.CyclicLR 和 torch.optim.lr_scheduler.OneCycleLR 兩種方法實現(xiàn)該操作，請參閱相關(guān)文檔。
這兩個方法的一個缺點是引入了許多額外的超參數(shù)。這篇文章和倉庫對如何查找好的超參數(shù)（包括上文提及的學(xué)習(xí)率）提供了詳細(xì)概述和實現(xiàn)。
至于為什么要這樣做？現(xiàn)今并不完全清楚，但一個可能的解釋是：定期提高學(xué)習(xí)率有助于更快越過損失鞍點。
在使用 torch.utils.data.DataLoader時，令 num_workers > 0，而不是默認(rèn)值 0，同時設(shè)置 pin_memory=True，而不是默認(rèn)值 False。至于為什么這么做，這篇文章會給你答案。
根據(jù)上述方法，Szymon Micacz 在四個 worker 和頁鎖定內(nèi)存的情況下，在單個epoch中實現(xiàn)了 2 倍加速。
根據(jù)經(jīng)驗，一般將進(jìn)程數(shù)量設(shè)置為可用 GPU 數(shù)量的四倍，大于或小于這個值都會降低訓(xùn)練速度。但是要注意，增加num_workers會增加 CPU 內(nèi)存消耗。
一直以來，人們對于調(diào)大batch沒有定論。一般來說，在GPU內(nèi)存允許的情況下增大batch將會增快訓(xùn)練速度，但同時還需要調(diào)整學(xué)習(xí)率等其他超參數(shù)。根據(jù)經(jīng)驗，batch大小加倍時，學(xué)習(xí)率也相應(yīng)加倍。
OpenAI 的論文表明不同的batch大小收斂周期不同。Daniel Huynh用不同的batch大小進(jìn)行了一些實驗（使用上述1Cycle 策略），實驗中他將 batch大小由64增加到512，實現(xiàn)了4倍加速。
然而也要注意，較大的batch會降低模型泛化能力，反之亦然。
PyTorch1.6支持本地自動混合精度訓(xùn)練。與單精度 (FP32) 相比，一些運(yùn)算在不損失準(zhǔn)確率的情況下，使用半精度 (FP16)速度更快。AMP能夠自動決定應(yīng)該以哪種精度執(zhí)行哪種運(yùn)算，這樣既可以加快訓(xùn)練速度，又減少了內(nèi)存占用。
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() |
Huang及其同事在NVIDIA V100 GPU上對一些常用語言和視覺模型進(jìn)行了基準(zhǔn)測試，發(fā)現(xiàn)在FP32訓(xùn)練中使用AMP提高約2倍的訓(xùn)練速度，最高甚至達(dá)到5.5倍。
目前，只有CUDA支持上述方式，查看本文檔了解更多信息。
AdamW是由fast.ai提出的具有權(quán)重衰減（而非 L2 正則化）的Adam， PyTorch中通過torch.optim.AdamW實現(xiàn)。在誤差和訓(xùn)練時間上，AdamW都優(yōu)于Adam。查看此文章了解為什么權(quán)重衰減使得Adam產(chǎn)生更好效果。
Adam和AdamW都很適合前文提到的1Cycle策略。
此外，LARS和LAMB等其他優(yōu)化器也收到廣泛關(guān)注。
NVIDA的APEX對Adam等常見優(yōu)化器進(jìn)行優(yōu)化融合，相比PyTorch中的原始Adam，由于避免了GPU內(nèi)存之間的多次傳遞，訓(xùn)練速度提升約 5%。
如果你的模型架構(gòu)時固定的，同時輸入大小保持不變，那么設(shè)置torch.backends.cudnn.benchmark = True可能會提升模型速度（幫助文檔）。通過啟用cudNN自動調(diào)節(jié)器，可以在cudNN中對多種計算卷積的方法進(jìn)行基準(zhǔn)測試，然后選擇最快的方法。
至于提速效果，Szymon Migacz在前向卷積時提速70％，在同時向前和后向卷積時提升了27％。
注意，如果你想要根據(jù)上述方法最大化批大小，該自動調(diào)整可能會非常耗時。
通過tensor.cpu()可以將張量從GPU傳輸?shù)紺PU，反之使用tensor.cuda()，但這樣的數(shù)據(jù)轉(zhuǎn)化代價較高。 .item()和.numpy()的使用也是如此，建議使用.detach（）。
如果要創(chuàng)建新的張量，使用關(guān)鍵字參數(shù)device=torch.device（'cuda：0'）將其直接分配給GPU。
最好使用.to(non_blocking=True)傳輸數(shù)據(jù)，確保傳輸后沒有任何同步點即可。
另外Santosh Gupta的SpeedTorch也值得一試，盡管其加速與否尚不完全清除。
檢查點通過將計算保存到內(nèi)存來工作。檢查點在反向傳播算法過程中并不保存計算圖的中間激活，而是在反向傳播時重新計算，其可用于模型的任何部分。
具體來說，在前向傳播中，function以torch.no_grad()方式運(yùn)行，不存儲任何中間激活。相反，前向傳遞將保存輸入元組和function參數(shù)。在反向傳播時，檢索保存的輸入和function，并再次對function進(jìn)行正向傳播，記錄中間激活，并使用這些激活值計算梯度。
因此，對于特定的批處理大小，這可能會稍微增加運(yùn)行時間，但會顯著減少內(nèi)存消耗。反過來，你可以進(jìn)一步增加批處理大小，從而更好地利用GPU。
雖然檢查點可以通過torch.utils.checkpoint方便實現(xiàn)，但仍需要里哦阿姐其思想與本質(zhì)。Priya Goyal的教程很清晰的演示了檢查點的一些關(guān)鍵思想，推薦閱讀。
增加批處理大小的另一種方法是在調(diào)用Optimizer.step()之對多個.backward()傳遞梯度進(jìn)行累積。
根據(jù)Hugging Face的Thomas Wolf發(fā)表的文章，可以按以下方式實現(xiàn)梯度累積：
| model.zero_grad() # Reset gradients tensors for 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 accumulated |
該方法主要是為了規(guī)避GPU內(nèi)存的限制，但對其他.backward()循環(huán)之間的取舍我并不清楚。fastai論壇上的討論似乎表明它實際上是可以加速訓(xùn)練的，因此值得一試。詳情查看GitHub托管的rawgradient_accumulation.py。
通過分布式訓(xùn)練加快模型速度的一種簡單的方法是使用torch.nn.DistributedDataParallel而不是torch.nn.DataParallel。這樣，每個GPU將由專用的CPU內(nèi)核驅(qū)動，從而避免了DataParallel的GIL問題。
強(qiáng)烈推薦閱讀分布式訓(xùn)練相關(guān)文檔了解更多信息：
PyTorch Distributed Overview — PyTorch Tutorials 1.7.0 documentation
設(shè)置.zero_grad(set_to_none=True)而不是.zero_grad()。
這樣內(nèi)存分配器處理梯度而不是主動將其設(shè)置為0，這會產(chǎn)生該文檔所示的適度加速，但不要抱有過大期望。
注意，這樣做不會有任何副作用！閱讀文檔查看更多信息。
12.使用.as_tensor()而不是.tensor()
torch.tensor()本質(zhì)是復(fù)制數(shù)據(jù)，因此，如果要轉(zhuǎn)換numpy數(shù)組，使用torch.as_tensor()或torch.from_numpy()可以避免復(fù)制數(shù)據(jù)。
Pytorch提供了許多調(diào)試工具，例如autograd.profiler, autograd.grad_check和autograd.anomaly_detection。使用時一定要謹(jǐn)慎，這些調(diào)試工具顯然會影響訓(xùn)練速度，因此在不需要時將其關(guān)閉。
為了避免RNN中的梯度爆炸，使用梯度裁剪gradient = min(gradient, threshold)可以起到加速收斂作用，這一方法已得到理論和實驗的支持。
Hugging Face的Transformer提供了將梯度裁剪和AMP等其他方法有效結(jié)合的清晰示例。
在PyTorch中，也可使用torch.nn.utils.clip_grad_norm_（文檔查閱）完成此操作。
雖然我尚不完全清楚哪種模型可以從梯度裁剪中受益，但毫無疑問的是，對于RNN、基于Transformer和ResNets結(jié)構(gòu)的一系列優(yōu)化器來說，該方法顯然是起到一定作用的。
在BatchNormalization層之前關(guān)閉之前層的偏差時一種簡單有效的方法。對于二維卷積層，可以通過將bias關(guān)鍵字設(shè)置為False實現(xiàn)，即torch.nn.Conv2d(..., bias=False, ...)。閱讀該文檔了解其原理。
與其他方法相比，該方法的速度提升是有的。
在模型驗證時令torch.no_grad()
也許你已經(jīng)在這樣做了，但還是要仔細(xì)檢查，反復(fù)確認(rèn)：
是否規(guī)范化輸入？
是否規(guī)范化批處理？
點擊查看這樣做的原因。
如果要執(zhí)行相鄰逐點操作，可以使用PyTorch JIT將它們組合成一個FusionGroup，然后在單內(nèi)核上啟動，而不是像默認(rèn)情況那樣在多個內(nèi)核上啟動，同時還可以保存一些內(nèi)存進(jìn)行讀寫。
Szymon Migacz展示了如何使用@torch.jit.script裝飾器融合GELU操作融合，如下：
| @torch.jit.scriptdef fused_gelu(x): return x * 0.5 * (1.0 + torch.erf(x / 1.41421)) |
相比于未融合版本，融合這些操作可以使fused_gelu的執(zhí)行速度提高5倍。
查閱此文章獲取更多使用Torchscript加速RNN的示例。
當(dāng)然，你還可以在Reddit上與u/Patient_Atmosphere45交流討論。
本文許多技巧參考自Szymon Migacz的演講及PyTorch文檔。
PyTorch Lightning的作者William Falcon在這兩篇文章種介紹了關(guān)于加快訓(xùn)練的內(nèi)容。同時，PyTorch Lightning已集成以上一些技巧與方法。
Hugging Face的作者Thomas Wolf也寫了一系列文章介紹深度學(xué)習(xí)的加速-尤其是語言模型。
Sylvain Gugger和Jeremy Howard寫了很多關(guān)于學(xué)習(xí)率和AdamW的文章。
感謝Ben Hahn，Kevin Klein和Robin Vaaler對本文撰寫提供的幫助！
AI研習(xí)社是AI學(xué)術(shù)青年和AI開發(fā)者技術(shù)交流的在線社區(qū)。我們與高校、學(xué)術(shù)機(jī)構(gòu)和產(chǎn)業(yè)界合作，通過提供學(xué)習(xí)、實戰(zhàn)和求職服務(wù)，為AI學(xué)術(shù)青年和開發(fā)者的交流互助和職業(yè)發(fā)展打造一站式平臺，致力成為中國最大的科技創(chuàng)新人才聚集地。
如果，你也是位熱愛分享的AI愛好者。歡迎與譯站一起，學(xué)習(xí)新知，分享成長。
雷峰網(wǎng)版權(quán)文章，未經(jīng)授權(quán)禁止轉(zhuǎn)載。詳情見轉(zhuǎn)載須知。
