PyTorch | 優(yōu)化神經(jīng)網(wǎng)絡(luò)訓(xùn)練的十七種方法
本文介紹在使用 PyTorch 訓(xùn)練深度模型時(shí)最省力、最有效的 17 種方法。該文所提方法，都是假設(shè)你在 GPU 環(huán)境下訓(xùn)練模型。具體內(nèi)容如下。
01. 考慮換一種學(xué)習(xí)率 schedule
學(xué)習(xí)率 schedule 的選擇對(duì)模型的收斂速度和泛化能力有很大的影響。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 對(duì)其進(jìn)行了推廣。下圖是 1Cycle 學(xué)習(xí)率 schedule 的圖示：
Sylvain 寫(xiě)到：1Cycle 包括兩個(gè)等長(zhǎng)的步幅，一個(gè)步幅是從較低的學(xué)習(xí)率到較高的學(xué)習(xí)率，另一個(gè)是回到最低水平。最大值來(lái)自學(xué)習(xí)率查找器選取的值，較小的值可以低十倍。然后，這個(gè)周期的長(zhǎng)度應(yīng)該略小于總的 epochs 數(shù)，并且，在訓(xùn)練的最后階段，我們應(yīng)該允許學(xué)習(xí)率比最小值小幾個(gè)數(shù)量級(jí)。
與傳統(tǒng)的學(xué)習(xí)率 schedule 相比，在最好的情況下，該 schedule 實(shí)現(xiàn)了巨大的加速（Smith 稱之為超級(jí)收斂）。例如，使用 1Cycle 策略在 ImageNet 數(shù)據(jù)集上訓(xùn)練 ResNet-56，訓(xùn)練迭代次數(shù)減少為原來(lái)的 1/10，但模型性能仍能比肩原論文中的水平。在常見(jiàn)的體系架構(gòu)和優(yōu)化器中，這種 schedule 似乎表現(xiàn)得很好。
Pytorch 已經(jīng)實(shí)現(xiàn)了這兩種方法：「torch.optim.lr_scheduler.CyclicLR」和「torch.optim.lr_scheduler.OneCycleLR」。
參考文檔：https://pytorch.org/docs/stable/optim.html
02. 在 DataLoader 中使用多個(gè) worker 和頁(yè)鎖定內(nèi)存
當(dāng)使用 torch.utils.data.DataLoader 時(shí)，設(shè)置 num_workers > 0，而不是默認(rèn)值 0，同時(shí)設(shè)置 pin_memory=True，而不是默認(rèn)值 False。
參考文檔：https://pytorch.org/docs/stable/data.html
來(lái)自 NVIDIA 的高級(jí) CUDA 深度學(xué)習(xí)算法軟件工程師 Szymon Micacz 就曾使用四個(gè) worker 和頁(yè)鎖定內(nèi)存（pinned memory）在單個(gè) epoch 中實(shí)現(xiàn)了 2 倍的加速。人們選擇 worker 數(shù)量的經(jīng)驗(yàn)法則是將其設(shè)置為可用 GPU 數(shù)量的四倍，大于或小于這個(gè)數(shù)都會(huì)降低訓(xùn)練速度。請(qǐng)注意，增加 num_workers 將增加 CPU 內(nèi)存消耗。
03. 把 batch 調(diào)到最大
把 batch 調(diào)到最大是一個(gè)頗有爭(zhēng)議的觀點(diǎn)。一般來(lái)說(shuō)，如果在 GPU 內(nèi)存允許的范圍內(nèi)將 batch 調(diào)到最大，你的訓(xùn)練速度會(huì)更快。但是，你也必須調(diào)整其他超參數(shù)，比如學(xué)習(xí)率。一個(gè)比較好用的經(jīng)驗(yàn)是，batch 大小加倍時(shí)，學(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 大小進(jìn)行了一些實(shí)驗(yàn)（也使用上面討論的 1Cycle 策略）。最終，他將 batch 大小由 64 增加到 512，實(shí)現(xiàn)了 4 倍的加速。
然而，使用大 batch 的不足是，這可能導(dǎo)致解決方案的泛化能力比使用小 batch 的差。
04. 使用自動(dòng)混合精度（AMP）
PyTorch 1.6 版本包括對(duì) PyTorch 的自動(dòng)混合精度訓(xùn)練的本地實(shí)現(xiàn)。這里想說(shuō)的是，與單精度 (FP32) 相比，某些運(yùn)算在半精度 (FP16) 下運(yùn)行更快，而不會(huì)損失準(zhǔn)確率。AMP 會(huì)自動(dòng)決定應(yīng)該以哪種精度執(zhí)行哪種運(yùn)算。這樣既可以加快訓(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()
05. 考慮使用另一種優(yōu)化器
AdamW 是由 fast.ai 推廣的一種具有權(quán)重衰減（而不是 L2 正則化）的 Adam，在 PyTorch 中以 torch.optim.AdamW 實(shí)現(xiàn)。AdamW 似乎在誤差和訓(xùn)練時(shí)間上都一直優(yōu)于 Adam。
Adam 和 AdamW 都能與上面提到的 1Cycle 策略很好地搭配。
目前，還有一些非本地優(yōu)化器也引起了很大的關(guān)注，最突出的是 LARS 和 LAMB。NVIDA 的 APEX 實(shí)現(xiàn)了一些常見(jiàn)優(yōu)化器的融合版本，比如 Adam。與 PyTorch 中的 Adam 實(shí)現(xiàn)相比，這種實(shí)現(xiàn)避免了與 GPU 內(nèi)存之間的多次傳遞，速度提高了 5%。
06. cudNN 基準(zhǔn)
如果你的模型架構(gòu)保持不變、輸入大小保持不變，設(shè)置 torch.backends.cudnn.benchmark = True。
07. 小心 CPU 和 GPU 之間頻繁的數(shù)據(jù)傳輸
當(dāng)頻繁地使用 tensor.cpu() 將張量從 GPU 轉(zhuǎn)到 CPU（或使用 tensor.cuda() 將張量從 CPU 轉(zhuǎn)到 GPU）時(shí)，代價(jià)是非常昂貴的。item() 和 .numpy() 也是一樣可以使用. detach() 代替。
如果你創(chuàng)建了一個(gè)新的張量，可以使用關(guān)鍵字參數(shù) device=torch.device( cuda:0 ) 將其分配給 GPU。
如果你需要傳輸數(shù)據(jù)，可以使用. to(non_blocking=True)，只要在傳輸之后沒(méi)有同步點(diǎn)。
08. 使用梯度 / 激活 checkpointing
Checkpointing 的工作原理是用計(jì)算換內(nèi)存，并不存儲(chǔ)整個(gè)計(jì)算圖的所有中間激活用于 backward pass，而是重新計(jì)算這些激活。我們可以將其應(yīng)用于模型的任何部分。
具體來(lái)說(shuō)，在 forward pass 中，function 會(huì)以 torch.no_grad() 方式運(yùn)行，不存儲(chǔ)中間激活。相反的是， forward pass 中會(huì)保存輸入元組以及 function 參數(shù)。在 backward pass 中，輸入和 function 會(huì)被檢索，并再次在 function 上計(jì)算 forward pass。然后跟蹤中間激活，使用這些激活值計(jì)算梯度。
因此，雖然這可能會(huì)略微增加給定 batch 大小的運(yùn)行時(shí)間，但會(huì)顯著減少內(nèi)存占用。這反過(guò)來(lái)又將允許進(jìn)一步增加所使用的 batch 大小，從而提高 GPU 的利用率。
盡管 checkpointing 以 torch.utils.checkpoint 方式實(shí)現(xiàn)，但仍需要一些思考和努力來(lái)正確地實(shí)現(xiàn)。Priya Goyal 寫(xiě)了一個(gè)很好的教程來(lái)介紹 checkpointing 關(guān)鍵方面。
Priya Goyal 教程地址：
09. 使用梯度積累
增加 batch 大小的另一種方法是在調(diào)用 optimizer.step() 之前在多個(gè). backward() 傳遞中累積梯度。
Hugging Face 的 Thomas Wolf 的文章《Training Neural Nets on Larger Batches: Practical Tips for 1-GPU, Multi-GPU & Distributed setups》介紹了如何使用梯度累積。梯度累積可以通過(guò)如下方式實(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
- lossloss = 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 accumulate
這個(gè)方法主要是為了規(guī)避 GPU 內(nèi)存的限制而開(kāi)發(fā)的。
10. 使用分布式數(shù)據(jù)并行進(jìn)行多 GPU 訓(xùn)練
加速分布式訓(xùn)練可能有很多方法，但是簡(jiǎn)單的方法是使用 torch.nn.DistributedDataParallel 而不是 torch.nn.DataParallel。這樣一來(lái)，每個(gè) GPU 將由一個(gè)專用的 CPU 核心驅(qū)動(dòng)，避免了 DataParallel 的 GIL 問(wèn)題。
分布式訓(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)存分配器處理梯度，而不是將它們?cè)O(shè)置為 0。正如文檔中所說(shuō)，將梯度設(shè)置為 None 會(huì)產(chǎn)生適度的加速，但不要期待奇跡出現(xiàn)。注意，這樣做也有缺點(diǎn)，詳細(xì)信息請(qǐng)查看文檔。
文檔地址：https://pytorch.org/docs/stable/optim.html
12. 使用. as_tensor() 而不是. tensor()
torch.tensor() 總是會(huì)復(fù)制數(shù)據(jù)。如果你要轉(zhuǎn)換一個(gè) numpy 數(shù)組，使用 torch.as_tensor() 或 torch.from_numpy() 來(lái)避免復(fù)制數(shù)據(jù)。
13. 必要時(shí)打開(kāi)調(diào)試工具
PyTorch 提供了很多調(diào)試工具，例如 autograd.profiler、autograd.grad_check、autograd.anomaly_detection。請(qǐng)確保當(dāng)你需要調(diào)試時(shí)再打開(kāi)調(diào)試器，不需要時(shí)要及時(shí)關(guān)掉，因?yàn)檎{(diào)試器會(huì)降低你的訓(xùn)練速度。
14. 使用梯度裁剪
關(guān)于避免 RNN 中的梯度爆炸的問(wèn)題，已經(jīng)有一些實(shí)驗(yàn)和理論證實(shí)，梯度裁剪（gradient = min(gradient, threshold)）可以加速收斂。HuggingFace 的 Transformer 實(shí)現(xiàn)就是一個(gè)非常清晰的例子，說(shuō)明了如何使用梯度裁剪。本文中提到的其他一些方法，如 AMP 也可以用。
在 PyTorch 中可以使用 torch.nn.utils.clip_grad_norm_來(lái)實(shí)現(xiàn)。
15. 在 BatchNorm 之前關(guān)閉 bias
在開(kāi)始 BatchNormalization 層之前關(guān)閉 bias 層。對(duì)于一個(gè) 2-D 卷積層，可以將 bias 關(guān)鍵字設(shè)置為 False：torch.nn.Conv2d(..., bias=False, ...)。
16. 在驗(yàn)證期間關(guān)閉梯度計(jì)算
在驗(yàn)證期間關(guān)閉梯度計(jì)算，設(shè)置：torch.no_grad() 。
17. 使用輸入和 batch 歸一化
要再三檢查一下輸入是否歸一化？是否使用了 batch 歸一化？
