如果你正在 pytorch 中訓(xùn)練深度學(xué)習(xí)模型，那么如何能夠加快模型訓(xùn)練速度呢?

在本文中，我會(huì)介紹一些改動(dòng)最小、影響最大的在pytorch中加速深度學(xué)習(xí)模型的方法。對(duì)于每種方法，我會(huì)對(duì)其思路進(jìn)行簡(jiǎn)要介紹，然后預(yù)估提升速度并討論其限制。我會(huì)把我認(rèn)為重要的部分強(qiáng)調(diào)介紹，并在每個(gè)部分展示一些實(shí)例。接下來我將假設(shè)你正在使用GPU訓(xùn)練模型，這些方法基本不需要導(dǎo)入其他的庫，只需要再pytorch內(nèi)進(jìn)行更改即可。

[[378697]]

以下是我根據(jù)預(yù)估的加速效果對(duì)不同方法的排序：

1. 考慮使用其他的學(xué)習(xí)率調(diào)整計(jì)劃
2. 在DataLoader中使用多個(gè)輔助進(jìn)程并頁鎖定內(nèi)存
3. 最大化batch大小
4. 使用自動(dòng)混合精度AMP
5. 考慮不同的優(yōu)化器
6. 打開cudNN基準(zhǔn)
7. 當(dāng)心CPU與GPU之間的數(shù)據(jù)傳輸
8. 使用梯度/激活檢查點(diǎn)
9. 使用梯度累積
10. 多GPU分布式訓(xùn)練
11. 將梯度設(shè)置為None而不是0
12. 使用.as_tensor()而不是.tensor()
13. 只在需要的時(shí)候打開debugging模式
14. 使用梯度裁剪
15. 在BatchNorm之前忽略偏差
16. 驗(yàn)證時(shí)關(guān)閉梯度計(jì)算
17. 規(guī)范化輸入和批處理

1. 考慮使用其他的學(xué)習(xí)率調(diào)整計(jì)劃

在訓(xùn)練中使用的學(xué)習(xí)率調(diào)整計(jì)劃會(huì)極大影響收斂速率以及模型泛化能力。

Leslie N. Smith 提出了循環(huán)學(xué)習(xí)率和1Cycle 學(xué)習(xí)率方法，然后由 fast.ai 的 Jeremy Howard 和 Sylvain Gugger 推廣了?？偟膩碚f，1Cycle 學(xué)習(xí)速率方法如下圖所示：

在最好的情況下，與傳統(tǒng)的學(xué)習(xí)率策略相比，這種策略可以實(shí)現(xiàn)巨大的加速—— Smith稱之為“超級(jí)收斂”。例如，使用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 兩種方法實(shí)現(xiàn)該操作，請(qǐng)參閱相關(guān)文檔。

這兩個(gè)方法的一個(gè)缺點(diǎn)是引入了許多額外的超參數(shù)。這篇文章和倉庫對(duì)如何查找好的超參數(shù)(包括上文提及的學(xué)習(xí)率)提供了詳細(xì)概述和實(shí)現(xiàn)。

至于為什么要這樣做?現(xiàn)今并不完全清楚，但一個(gè)可能的解釋是：定期提高學(xué)習(xí)率有助于更快越過損失鞍點(diǎn)。

2. 在DataLoader中使用多個(gè)輔助進(jìn)程并頁鎖定內(nèi)存

在使用 torch.utils.data.DataLoader時(shí)，令 num_workers > 0，而不是默認(rèn)值 0，同時(shí)設(shè)置 pin_memory=True，而不是默認(rèn)值 False。至于為什么這么做，這篇文章會(huì)給你答案。

根據(jù)上述方法，Szymon Micacz 在四個(gè) worker 和頁鎖定內(nèi)存的情況下，在單個(gè)epoch中實(shí)現(xiàn)了 2 倍加速。

根據(jù)經(jīng)驗(yàn)，一般將進(jìn)程數(shù)量設(shè)置為可用 GPU 數(shù)量的四倍，大于或小于這個(gè)值都會(huì)降低訓(xùn)練速度。但是要注意，增加num_workers會(huì)增加 CPU 內(nèi)存消耗。

3.最大化batch大小

一直以來，人們對(duì)于調(diào)大batch沒有定論。一般來說，在GPU內(nèi)存允許的情況下增大batch將會(huì)增快訓(xùn)練速度，但同時(shí)還需要調(diào)整學(xué)習(xí)率等其他超參數(shù)。根據(jù)經(jīng)驗(yàn)，batch大小加倍時(shí)，學(xué)習(xí)率也相應(yīng)加倍。

OpenAI 的論文表明不同的batch大小收斂周期不同。Daniel Huynh用不同的batch大小進(jìn)行了一些實(shí)驗(yàn)(使用上述1Cycle 策略)，實(shí)驗(yàn)中他將 batch大小由64增加到512，實(shí)現(xiàn)了4倍加速。

然而也要注意，較大的batch會(huì)降低模型泛化能力，反之亦然。

4. 使用自動(dòng)混合精度AMP

PyTorch1.6支持本地自動(dòng)混合精度訓(xùn)練。與單精度 (FP32) 相比，一些運(yùn)算在不損失準(zhǔn)確率的情況下，使用半精度 (FP16)速度更快。AMP能夠自動(dòng)決定應(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上對(duì)一些常用語言和視覺模型進(jìn)行了基準(zhǔn)測(cè)試，發(fā)現(xiàn)在FP32訓(xùn)練中使用AMP提高約2倍的訓(xùn)練速度，最高甚至達(dá)到5.5倍。

目前，只有CUDA支持上述方式，查看本文檔了解更多信息。

5. 考慮不同的優(yōu)化器

AdamW是由fast.ai提出的具有權(quán)重衰減(而非 L2 正則化)的Adam， PyTorch中通過torch.optim.AdamW實(shí)現(xiàn)。在誤差和訓(xùn)練時(shí)間上，AdamW都優(yōu)于Adam。查看此文章了解為什么權(quán)重衰減使得Adam產(chǎn)生更好效果。

Adam和AdamW都很適合前文提到的1Cycle策略。

此外，LARS和LAMB等其他優(yōu)化器也收到廣泛關(guān)注。

NVIDA的APEX對(duì)Adam等常見優(yōu)化器進(jìn)行優(yōu)化融合，相比PyTorch中的原始Adam，由于避免了GPU內(nèi)存之間的多次傳遞，訓(xùn)練速度提升約 5%。

6. 打開cudNN基準(zhǔn)

如果你的模型架構(gòu)時(shí)固定的，同時(shí)輸入大小保持不變，那么設(shè)置torch.backends.cudnn.benchmark = True可能會(huì)提升模型速度(幫助文檔)。通過啟用cudNN自動(dòng)調(diào)節(jié)器，可以在cudNN中對(duì)多種計(jì)算卷積的方法進(jìn)行基準(zhǔn)測(cè)試，然后選擇最快的方法。

至于提速效果，Szymon Migacz在前向卷積時(shí)提速70%，在同時(shí)向前和后向卷積時(shí)提升了27%。

注意，如果你想要根據(jù)上述方法最大化批大小，該自動(dòng)調(diào)整可能會(huì)非常耗時(shí)。

7. 當(dāng)心CPU與GPU之間的數(shù)據(jù)傳輸

通過tensor.cpu()可以將張量從GPU傳輸?shù)紺PU，反之使用tensor.cuda()，但這樣的數(shù)據(jù)轉(zhuǎn)化代價(jià)較高。 .item()和.numpy()的使用也是如此，建議使用.detach()。

如果要?jiǎng)?chuàng)建新的張量，使用關(guān)鍵字參數(shù)device=torch.device('cuda：0')將其直接分配給GPU。

最好使用.to(non_blocking=True)傳輸數(shù)據(jù)，確保傳輸后沒有任何同步點(diǎn)即可。

另外Santosh Gupta的SpeedTorch也值得一試，盡管其加速與否尚不完全清除。

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

檢查點(diǎn)通過將計(jì)算保存到內(nèi)存來工作。檢查點(diǎn)在反向傳播算法過程中并不保存計(jì)算圖的中間激活，而是在反向傳播時(shí)重新計(jì)算，其可用于模型的任何部分。

具體來說，在前向傳播中，function以torch.no_grad()方式運(yùn)行，不存儲(chǔ)任何中間激活。相反，前向傳遞將保存輸入元組和function參數(shù)。在反向傳播時(shí)，檢索保存的輸入和function，并再次對(duì)function進(jìn)行正向傳播，記錄中間激活，并使用這些激活值計(jì)算梯度。

因此，對(duì)于特定的批處理大小，這可能會(huì)稍微增加運(yùn)行時(shí)間，但會(huì)顯著減少內(nèi)存消耗。反過來，你可以進(jìn)一步增加批處理大小，從而更好地利用GPU。

雖然檢查點(diǎn)可以通過torch.utils.checkpoint方便實(shí)現(xiàn)，但仍需要里哦阿姐其思想與本質(zhì)。Priya Goyal的教程很清晰的演示了檢查點(diǎn)的一些關(guān)鍵思想，推薦閱讀。

9.使用梯度累積

增加批處理大小的另一種方法是在調(diào)用Optimizer.step()之對(duì)多個(gè).backward()傳遞梯度進(jìn)行累積。

根據(jù)Hugging Face的Thomas Wolf發(fā)表的文章，可以按以下方式實(shí)現(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)存的限制，但對(duì)其他.backward()循環(huán)之間的取舍我并不清楚。fastai論壇上的討論似乎表明它實(shí)際上是可以加速訓(xùn)練的，因此值得一試。詳情查看GitHub托管的rawgradient_accumulation.py。

10.多GPU分布式訓(xùn)練

通過分布式訓(xùn)練加快模型速度的一種簡(jiǎn)單的方法是使用torch.nn.DistributedDataParallel而不是torch.nn.DataParallel。這樣，每個(gè)GPU將由專用的CPU內(nèi)核驅(qū)動(dòng)，從而避免了DataParallel的GIL問題。

強(qiáng)烈推薦閱讀分布式訓(xùn)練相關(guān)文檔了解更多信息：

PyTorch Distributed Overview — PyTorch Tutorials 1.7.0 documentation 

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

設(shè)置.zero_grad(set_to_none=True)而不是.zero_grad()。

這樣內(nèi)存分配器處理梯度而不是主動(dòng)將其設(shè)置為0，這會(huì)產(chǎn)生該文檔所示的適度加速，但不要抱有過大期望。

注意，這樣做不會(huì)有任何副作用!閱讀文檔查看更多信息。

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

torch.tensor()本質(zhì)是復(fù)制數(shù)據(jù)，因此，如果要轉(zhuǎn)換numpy數(shù)組，使用torch.as_tensor()或torch.from_numpy()可以避免復(fù)制數(shù)據(jù)。

13.只在需要的時(shí)候打開debugging模式

Pytorch提供了許多調(diào)試工具，例如autograd.profiler, autograd.grad_check和autograd.anomaly_detection。使用時(shí)一定要謹(jǐn)慎，這些調(diào)試工具顯然會(huì)影響訓(xùn)練速度，因此在不需要時(shí)將其關(guān)閉。

14.使用梯度裁剪

為了避免RNN中的梯度爆炸，使用梯度裁剪gradient = min(gradient, threshold)可以起到加速收斂作用，這一方法已得到理論和實(shí)驗(yàn)的支持。

Hugging Face的Transformer提供了將梯度裁剪和AMP等其他方法有效結(jié)合的清晰示例。

在PyTorch中，也可使用torch.nn.utils.clip_grad_norm_(文檔查閱)完成此操作。

雖然我尚不完全清楚哪種模型可以從梯度裁剪中受益，但毫無疑問的是，對(duì)于RNN、基于Transformer和ResNets結(jié)構(gòu)的一系列優(yōu)化器來說，該方法顯然是起到一定作用的。

15.在BatchNorm之前忽略偏差

在BatchNormalization層之前關(guān)閉之前層的偏差時(shí)一種簡(jiǎn)單有效的方法。對(duì)于二維卷積層，可以通過將bias關(guān)鍵字設(shè)置為False實(shí)現(xiàn)，即torch.nn.Conv2d(..., bias=False, ...)。閱讀該文檔了解其原理。

與其他方法相比，該方法的速度提升是有的。

16. 驗(yàn)證時(shí)關(guān)閉梯度計(jì)算

在模型驗(yàn)證時(shí)令torch.no_grad()

17. 規(guī)范化輸入和批處理

也許你已經(jīng)在這樣做了，但還是要仔細(xì)檢查，反復(fù)確認(rèn)：

• 是否規(guī)范化輸入?
• 是否規(guī)范化批處理?

其他技巧：使用JIT實(shí)現(xiàn)逐點(diǎn)融合

如果要執(zhí)行相鄰逐點(diǎn)操作，可以使用PyTorch JIT將它們組合成一個(gè)FusionGroup，然后在單內(nèi)核上啟動(dòng)，而不是像默認(rèn)情況那樣在多個(gè)內(nèi)核上啟動(dòng)，同時(shí)還可以保存一些內(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倍。

本文轉(zhuǎn)自雷鋒網(wǎng)，如需轉(zhuǎn)載請(qǐng)至雷鋒網(wǎng)官網(wǎng)申請(qǐng)授權(quán)。