比如說，你正在PyTorch中訓(xùn)練一個深度學(xué)習(xí)模型。你能做些什么讓你的訓(xùn)練更快結(jié)束？
在這篇文章中，我將概述一些在PyTorch中加速深度學(xué)習(xí)模型訓(xùn)練時改動最小，影響最大的方法。對于每種方法，我會簡要總結(jié)其思想，并估算預(yù)期的加速度，并討論一些限制。我將著重于傳達最重要的部分，并為每個部分給出額外的一些資源。大多數(shù)情況下，我會專注于可以直接在PyTorch中進行的更改，而不需要引入額外的庫，并且我將假設(shè)你正在使用GPU訓(xùn)練模型。
1. 考慮使用另外一種學(xué)習(xí)率策略
你選擇的學(xué)習(xí)率對收斂速度以及模型的泛化性能有很大的影響。
循環(huán)學(xué)習(xí)率和1Cycle學(xué)習(xí)率策略都是Leslie N. Smith提出的方法，然后由推廣。本質(zhì)上，1Cycle學(xué)習(xí)率策略看起來像這樣：
Sylvain寫道：
[1cycle由兩個相同長度的步驟組成，一個是從較低的學(xué)習(xí)率到較高的學(xué)習(xí)率，另一個步驟是回到最低的學(xué)習(xí)速率。最大值應(yīng)該是使用Learning Rate Finder選擇的值，較低的值可以低十倍。然后，這個周期的長度應(yīng)該略小于epochs的總數(shù)，并且，在訓(xùn)練的最后一部分，我們應(yīng)該允許學(xué)習(xí)率減少超過最小值幾個數(shù)量級。
在最好的情況下，與傳統(tǒng)的學(xué)習(xí)率策略相比，這種策略可以實現(xiàn)巨大的加速 —— Smith稱之為“超級收斂”。例如，使用1Cycle策略，在ImageNet上減少了ResNet-56訓(xùn)練迭代數(shù)的10倍，就可以匹配原始論文的性能。該策略似乎在通用架構(gòu)和優(yōu)化器之間運行得很好。
PyTorch實現(xiàn)了這兩個方法，torch.optim.lr_scheduler.CyclicLR和torch.optim.lr_scheduler.OneCycleLR。
這兩個策略的一個缺點是它們引入了許多額外的超參數(shù)。為什么會這樣呢？這似乎并不完全清楚，但一個可能的解釋是，定期提高學(xué)習(xí)率有助于更快的穿越鞍點。
2. 在 DataLoader中使用多個workers和pinned memory
當使用torch.utils.data.DataLoader時，設(shè)置num_workers > 0，而不是等于0，設(shè)置pin_memory=True而不是默認值False。詳細解釋：https://pytorch.org/docs/stable/data.html。
Szymon Micacz通過使用4個workers和pinned memory，實現(xiàn)了單個訓(xùn)練epoch的2倍加速。
一個經(jīng)驗法則，選擇workers的數(shù)量設(shè)置為可用GPU數(shù)量的4倍，更大或更小的workers數(shù)量會變慢。
注意，增加num_workers會增加CPU內(nèi)存消耗。
3. 最大化batch size
這是一個頗有爭議的觀點。一般來說，然而，似乎使用GPU允許的最大的batch size可能會加速你的訓(xùn)練。注意，如果你修改了batch大小，你還必須調(diào)整其他超參數(shù)，例如學(xué)習(xí)率。這里的一個經(jīng)驗法則是，當你把batch數(shù)量翻倍時，學(xué)習(xí)率也要翻倍。
OpenAI有一篇很好的實證論文關(guān)于不同batch size需要的收斂步驟的數(shù)量。Daniel Huynh運行一些實驗用不同batch大小(使用上面所討論的1Cycle策略)，從batch size 64到512他實現(xiàn)了4倍的加速。
然而，使用大batch的缺點之一是，它們可能會導(dǎo)致泛化能力比使用小batch的模型差。
4. 使用自動混合精度
PyTorch 1.6的發(fā)行版包含了對PyTorch進行自動混合精度訓(xùn)練的本地實現(xiàn)。這里的主要思想是，與在所有地方都使用單精度(FP32)相比，某些操作可以在半精度(FP16)下運行得更快，而且不會損失精度。然后，AMP自動決定應(yīng)該以何種格式執(zhí)行何種操作。這允許更快的訓(xùn)練和更小的內(nèi)存占用。
AMP的使用看起來像這樣：
import torch
# Creates once at the beginning of training
scaler = 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()
在NVIDIA V100 GPU上對多個NLP和CV的benchmark進行測試，Huang和他的同事們發(fā)現(xiàn)使用AMP在FP32訓(xùn)練收益率常規(guī)大約2x，但最高可達5.5x。
目前，只有CUDA ops可以通過這種方式自動轉(zhuǎn)換。
5. 考慮使用另外的優(yōu)化器
AdamW是由推廣的具有權(quán)重衰減(而不是L2正則化)的Adam。現(xiàn)在可以在PyTorch中直接使用，torch.optim.AdamW。無論在誤差還是訓(xùn)練時間上，AdamW都比Adam表現(xiàn)更好。
Adam和AdamW都可以很好地使用上面描述的1Cycle策略。
還有一些自帶優(yōu)化器最近受到了很多關(guān)注，最著名的是LARS和LAMB。
NVIDA的APEX實現(xiàn)了許多常見優(yōu)化器的融合版本，如Adam。與Adam的PyTorch實現(xiàn)相比，這種實現(xiàn)避免了大量進出GPU內(nèi)存的操作，從而使速度提高了5%。
6. 開啟cudNN benchmarking
如果你的模型架構(gòu)保持不變，你的輸入大小保持不變，設(shè)置torch.backends.cudnn.benchmark = True可能是有益的。這使得cudNN能夠測試許多不同的卷積計算方法，然后使用最快的方法。
對于加速的預(yù)期有一個粗略的參考，Szymon Migacz達到70%的forward的加速以及27%的forward和backward的加速。
這里需要注意的是，如果你像上面提到的那樣將batch size最大化，那么這種自動調(diào)優(yōu)可能會變得非常緩慢。
7. 注意CPU和GPU之間頻繁的數(shù)據(jù)傳輸
小心使用tensor.cpu()和tensor.cuda()頻繁地將張量從GPU和CPU之間相互轉(zhuǎn)換。對于.item()和.numpy()也是一樣，用.detach()代替。
如果你正在創(chuàng)建一個新的張量，你也可以使用關(guān)鍵字參數(shù)device=torch.device('cuda:0')直接將它分配給你的GPU。
如果你確實需要傳輸數(shù)據(jù)，在傳輸后使用.to(non_blocking=True)可能會很有用，只要你沒有任何同步點。
如果你真的需要，你可以試試Santosh Gupta的SpeedTorch，雖然不是很確定在什么情況下可以加速。
8. 使用gradient/activation檢查點
直接引用文檔中的話：
檢查點的工作原理是用計算交換內(nèi)存，并不是存儲整個計算圖的所有中間激活用于向后計算，檢查點不保存中間的激活，而是在向后傳遞中重新計算它們。可以應(yīng)用于模型的任何部分。
具體來說，在向前傳遞中，function會以torch.no_grad()的方式運行，也就是說，不存儲中間激活。相反，正向傳遞保存輸入和function的參數(shù)。在向后傳遞中，將檢索保存的輸入和function，并再次根據(jù)function計算向前傳遞，然后跟蹤中間的激活，再使用這些激活值計算梯度。
因此，雖然這可能會略微增加給定batch大小的運行時間，但會顯著減少內(nèi)存占用。這反過來會允許你進一步增加你正在使用的batch大小，從而更好地利用GPU。
檢查點的pytorch實現(xiàn)為torch.utils.checkpoint，需要想點辦法才能實現(xiàn)的很好。
9. 使用梯度累加
增加batch大小的另一種方法是在調(diào)用optimizer.step()之前，在多個.backward()中累積梯度。
在Hugging Face的實現(xiàn)中，梯度累加可以實現(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
這個方法主要是為了避開GPU內(nèi)存限制。fastai論壇上的這個討論：https://forums./t/accumulating-gradients/33219/28似乎表明它實際上可以加速訓(xùn)練，所以可能值得一試。
10. 對于多個GPU使用分布式數(shù)據(jù)并行
對于分布式訓(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。正如在文檔中所說的那樣，這會導(dǎo)致產(chǎn)生一個適度的加速，所以不要期待任何奇跡。
注意，這樣做并不是沒有副作用的！關(guān)于這一點的詳細信息請查看文檔。
12. 使用.as_tensor() 而不是 .tensor()
torch.tensor() 會拷貝數(shù)據(jù)，如果你有一個numpy數(shù)組，你想轉(zhuǎn)為tensor，使用 torch.as_tensor() 或是 torch.from_numpy() 來避免拷貝數(shù)據(jù)。
13. 需要的時候打開調(diào)試工具
Pytorch提供了大量的有用的調(diào)試工具，如autograd.profiler，autograd.grad_check和autograd.anomaly_detection。在需要的時候使用它們，在不需要它們的時候關(guān)閉它們，因為它們會減慢你的訓(xùn)練。
14. 使用梯度剪裁
最初是用于RNNs避免爆炸梯度，有一些經(jīng)驗證據(jù)和一些理論支持認為剪裁梯度(粗略地說：gradient = min(gradient, threshold))可以加速收斂。Hugging Face的Transformer實現(xiàn)是關(guān)于如何使用梯度剪裁以及其他的一些方法如AMP的一個非常干凈的例子。
在PyTorch中，這可以通過使用torch.nn.utils.clip_grad_norm_實現(xiàn)。我并不完全清楚哪個模型從梯度裁剪中獲益多少，但它似乎對RNN、基于Transformer和ResNets架構(gòu)以及一系列不同的優(yōu)化器都非常有用。
15. 在BatchNorm之前不使用bias
這是一個非常簡單的方法：在BatchNormalization 層之前不使用bias。對于二維卷積層，可以將關(guān)鍵字bias設(shè)為False: torch.nn.Conv2d(..., bias=False, ...)。
你會保存一些參數(shù)，然而，與這里提到的其他一些方法相比，我對這個方法的加速期望相對較小。
16. 在驗證的時候關(guān)閉梯度計算
這個很直接：在驗證的時候使用 torch.no_grad() 。
17. 對輸入和batch使用歸一化
你可能已經(jīng)這么做了，但你可能想再檢查一下:
- 你是否在使用batch-normalization
來自評論的額外的技巧：使用 JIT融合point-wise的操作
如果你有point-wise的操作，你可以使用PyTorch JIT將它們合并成一個FusionGroup，這樣就可以在單個核上啟動，而不是像默認情況下那樣在多個核上啟動。你還可以節(jié)省一些內(nèi)存的讀寫。
Szymon Migacz展示了如何使用@torch.jit腳本裝飾器來融合GELU中的操作，例如：
@torch.jit.script
def fused_gelu(x):
return x * 0.5 * (1.0 + torch.erf(x / 1.41421))
在本例中，與未融合的版本相比，融合操作將導(dǎo)致fused_gelu的執(zhí)行速度提高5倍。
一些相關(guān)的資源
上面列出的許多技巧來自Szymon Migacz的談話，并發(fā)表在：https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html。
PyTorch Lightning的William Falcon有兩篇文章：
https:///9-tips-for-training-lightning-fast-neural-networks-in-pytorch-8e63a502f565
https:///7-tips-for-squeezing-maximum-performance-from-pytorch-ca4a40951259
其中有加速訓(xùn)練的技巧。PyTorch Lightning已經(jīng)處理了上面默認的一些點。
Hugging Face的Thomas Wolf有很多關(guān)于加速深度學(xué)習(xí)的有趣文章，其中特別關(guān)注語言模型。
Sylvain Gugger和Jeremy Howard也有一些文章：
關(guān)于學(xué)習(xí)率策略的：https://sgugger./the-1cycle-policy.html，
關(guān)于找最佳學(xué)習(xí)率的：https://sgugger./how-do-you-find-a-good-learning-rate.html
AdamW相關(guān)的：https://www./2018/07/02/adam-weight-decay/。
