掌握這 17 種方法，用最省力的方式，加速你的 Pytorch 深度學(xué)習(xí)訓(xùn)練。

近日，Reddit 上一個帖子熱度爆表。主題內(nèi)容是關(guān)于怎樣加速 PyTorch 訓(xùn)練。原文作者是來自蘇黎世聯(lián)邦理工學(xué)院的計算機科學(xué)碩士生 LORENZ KUHN，文章向我們介紹了在使用 PyTorch 訓(xùn)練深度模型時最省力、最有效的 17 種方法。

該文所提方法，都是假設(shè)你在 GPU 環(huán)境下訓(xùn)練模型。具體內(nèi)容如下。

17 種加速 PyTorch 訓(xùn)練的方法

1. 考慮換一種學(xué)習(xí)率 schedule

學(xué)習(xí)率 schedule 的選擇對模型的收斂速度和泛化能力有很大的影響。Leslie N. Smith 等人在論文《Cyclical Learning Rates for Training Neural Networks》、《Super-Convergence: Very Fast Training of Neural Networks Using Large Learning Rates 》中提出了周期性（Cyclical）學(xué)習(xí)率以及 1Cycle 學(xué)習(xí)率 schedule。之后，fast.ai 的 Jeremy Howard 和 Sylvain Gugger 對其進行了推廣。下圖是 1Cycle 學(xué)習(xí)率 schedule 的圖示：

Sylvain 寫到：1Cycle 包括兩個等長的步幅，一個步幅是從較低的學(xué)習(xí)率到較高的學(xué)習(xí)率，另一個是回到最低水平。最大值來自學(xué)習(xí)率查找器選取的值，較小的值可以低十倍。然后，這個周期的長度應(yīng)該略小于總的 epochs 數(shù)，并且，在訓(xùn)練的最后階段，我們應(yīng)該允許學(xué)習(xí)率比最小值小幾個數(shù)量級。

與傳統(tǒng)的學(xué)習(xí)率 schedule 相比，在最好的情況下，該 schedule 實現(xiàn)了巨大的加速（Smith 稱之為超級收斂）。例如，使用 1Cycle 策略在 ImageNet 數(shù)據(jù)集上訓(xùn)練 ResNet-56，訓(xùn)練迭代次數(shù)減少為原來的 1/10，但模型性能仍能比肩原論文中的水平。在常見的體系架構(gòu)和優(yōu)化器中，這種 schedule 似乎表現(xiàn)得很好。

Pytorch 已經(jīng)實現(xiàn)了這兩種方法：「torch.optim.lr_scheduler.CyclicLR」和「torch.optim.lr_scheduler.OneCycleLR」。

2. 在 DataLoader 中使用多個 worker 和頁鎖定內(nèi)存

當使用 torch.utils.data.DataLoader 時，設(shè)置 num_workers > 0，而不是默認值 0，同時設(shè)置 pin_memory=True，而不是默認值 False。

參考文檔：https://pytorch.org/docs/stable/data.html

來自 NVIDIA 的高級 CUDA 深度學(xué)習(xí)算法軟件工程師 Szymon Micacz 就曾使用四個 worker 和頁鎖定內(nèi)存（pinned memory）在單個 epoch 中實現(xiàn)了 2 倍的加速。人們選擇 worker 數(shù)量的經(jīng)驗法則是將其設(shè)置為可用 GPU 數(shù)量的四倍，大于或小于這個數(shù)都會降低訓(xùn)練速度。請注意，增加 num_workers 將增加 CPU 內(nèi)存消耗。

3. 把 batch 調(diào)到最大

把 batch 調(diào)到最大是一個頗有爭議的觀點。一般來說，如果在 GPU 內(nèi)存允許的范圍內(nèi)將 batch 調(diào)到最大，你的訓(xùn)練速度會更快。但是，你也必須調(diào)整其他超參數(shù)，比如學(xué)習(xí)率。一個比較好用的經(jīng)驗是，batch 大小加倍時，學(xué)習(xí)率也要加倍。

OpenAI 的論文《An Empirical Model of Large-Batch Training》很好地論證了不同的 batch 大小需要多少步才能收斂。在《How to get 4x speedup and better generalization using the right batch size》一文中，作者 Daniel Huynh 使用不同的 batch 大小進行了一些實驗（也使用上面討論的 1Cycle 策略）。最終，他將 batch 大小由 64 增加到 512，實現(xiàn)了 4 倍的加速。

然而，使用大 batch 的不足是，這可能導(dǎo)致解決方案的泛化能力比使用小 batch 的差。

4. 使用自動混合精度（AMP）

PyTorch 1.6 版本包括對 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()

5. 考慮使用另一種優(yōu)化器

AdamW 是由 fast.ai 推廣的一種具有權(quán)重衰減（而不是 L2 正則化）的 Adam，在 PyTorch 中以 torch.optim.AdamW 實現(xiàn)。AdamW 似乎在誤差和訓(xùn)練時間上都一直優(yōu)于 Adam。

Adam 和 AdamW 都能與上面提到的 1Cycle 策略很好地搭配。

目前，還有一些非本地優(yōu)化器也引起了很大的關(guān)注，最突出的是 LARS 和 LAMB。NVIDA 的 APEX 實現(xiàn)了一些常見優(yōu)化器的融合版本，比如 Adam。與 PyTorch 中的 Adam 實現(xiàn)相比，這種實現(xiàn)避免了與 GPU 內(nèi)存之間的多次傳遞，速度提高了 5%。

6. cudNN 基準

如果你的模型架構(gòu)保持不變、輸入大小保持不變，設(shè)置 torch.backends.cudnn.benchmark = True。

7. 小心 CPU 和 GPU 之間頻繁的數(shù)據(jù)傳輸

當頻繁地使用 tensor.cpu() 將張量從 GPU 轉(zhuǎn)到 CPU（或使用 tensor.cuda() 將張量從 CPU 轉(zhuǎn)到 GPU）時，代價是非常昂貴的。item() 和 .numpy() 也是一樣可以使用. detach() 代替。

如果你創(chuàng)建了一個新的張量，可以使用關(guān)鍵字參數(shù) device=torch.device('cuda:0') 將其分配給 GPU。

如果你需要傳輸數(shù)據(jù)，可以使用. to(non_blocking=True)，只要在傳輸之后沒有同步點。

8. 使用梯度 / 激活 checkpointing

Checkpointing 的工作原理是用計算換內(nèi)存，并不存儲整個計算圖的所有中間激活用于 backward pass，而是重新計算這些激活。我們可以將其應(yīng)用于模型的任何部分。

具體來說，在 forward pass 中，function 會以 torch.no_grad() 方式運行，不存儲中間激活。相反的是， forward pass 中會保存輸入元組以及 function 參數(shù)。在 backward pass 中，輸入和 function 會被檢索，并再次在 function 上計算 forward pass。然后跟蹤中間激活，使用這些激活值計算梯度。

因此，雖然這可能會略微增加給定 batch 大小的運行時間，但會顯著減少內(nèi)存占用。這反過來又將允許進一步增加所使用的 batch 大小，從而提高 GPU 的利用率。

盡管 checkpointing 以 torch.utils.checkpoint 方式實現(xiàn)，但仍需要一些思考和努力來正確地實現(xiàn)。Priya Goyal 寫了一個很好的教程來介紹 checkpointing 關(guān)鍵方面。

9. 使用梯度積累

增加 batch 大小的另一種方法是在調(diào)用 optimizer.step() 之前在多個. backward() 傳遞中累積梯度。

Hugging Face 的 Thomas Wolf 的文章《Training Neural Nets on Larger Batches: Practical Tips for 1-GPU, Multi-GPU & Distributed setups》介紹了如何使用梯度累積。梯度累積可以通過如下方式實現(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 accumulat

這個方法主要是為了規(guī)避 GPU 內(nèi)存的限制而開發(fā)的。

10. 使用分布式數(shù)據(jù)并行進行多 GPU 訓(xùn)練

加速分布式訓(xùn)練可能有很多方法，但是簡單的方法是使用 torch.nn.DistributedDataParallel 而不是 torch.nn.DataParallel。這樣一來，每個 GPU 將由一個專用的 CPU 核心驅(qū)動，避免了 DataParallel 的 GIL 問題。

分布式訓(xùn)練文檔地址：https://pytorch.org/tutorials/beginner/dist_overview.html

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

梯度設(shè)置為. zero_grad(set_to_none=True) 而不是 .zero_grad()。這樣做可以讓內(nèi)存分配器處理梯度，而不是將它們設(shè)置為 0。正如文檔中所說，將梯度設(shè)置為 None 會產(chǎn)生適度的加速，但不要期待奇跡出現(xiàn)。注意，這樣做也有缺點，詳細信息請查看文檔。

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

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

13. 必要時打開調(diào)試工具

PyTorch 提供了很多調(diào)試工具，例如 autograd.profiler、autograd.grad_check、autograd.anomaly_detection。請確保當你需要調(diào)試時再打開調(diào)試器，不需要時要及時關(guān)掉，因為調(diào)試器會降低你的訓(xùn)練速度。

14. 使用梯度裁剪

關(guān)于避免 RNN 中的梯度爆炸的問題，已經(jīng)有一些實驗和理論證實，梯度裁剪（gradient = min(gradient, threshold)）可以加速收斂。HuggingFace 的 Transformer 實現(xiàn)就是一個非常清晰的例子，說明了如何使用梯度裁剪。本文中提到的其他一些方法，如 AMP 也可以用。

在 PyTorch 中可以使用 torch.nn.utils.clip_grad_norm_來實現(xiàn)。

15. 在 BatchNorm 之前關(guān)閉 bias

在開始 BatchNormalization 層之前關(guān)閉 bias 層。對于一個 2-D 卷積層，可以將 bias 關(guān)鍵字設(shè)置為 False：torch.nn.Conv2d(..., bias=False, ...)。

16. 在驗證期間關(guān)閉梯度計算

在驗證期間關(guān)閉梯度計算，設(shè)置：torch.no_grad() 。

17. 使用輸入和 batch 歸一化

要再三檢查一下輸入是否歸一化？是否使用了 batch 歸一化？