|
隨著深度學(xué)習(xí)快速發(fā)展，同時伴隨著模型參數(shù)的爆炸式增長，對顯卡的顯存容量提出了越來越高的要求，如何在單卡小容量顯卡上面訓(xùn)練模型是一直以來大家關(guān)心的問題。 本文結(jié)合 MMCV 開源庫對一些常用的節(jié)省顯存策略進行簡要分析，并希望能夠?qū)θ绾渭床寮从玫氖褂?MMCV 中相應(yīng)節(jié)省顯存策略的用戶提供一個簡單的指引。 本文涉及到的 PyTorch 節(jié)省顯存的策略包括： - 混合精度訓(xùn)練 - 大 batch 訓(xùn)練或者稱為梯度累加 - gradient checkpointing 梯度檢查點  本文內(nèi)容 混合精度訓(xùn)練 大 Batch 訓(xùn)練（梯度累加） 梯度檢查點 實驗驗證 1. 混合精度訓(xùn)練 混合精度訓(xùn)練全稱為 Automatic Mixed Precision，簡稱為 AMP，也就是我們常說的 FP16。在前系列解讀中已經(jīng)詳細分析了 AMP 原理、源碼實現(xiàn)以及 MMCV 中如何一行代碼使用 AMP，具體鏈接見： PyTorch 源碼解讀之 torch.cuda.amp: 自動混合精度詳解： https://zhuanlan.zhihu.com/p/348554267 OpenMMLab 中混合精度訓(xùn)練 AMP 的正確打開方式： https://zhuanlan.zhihu.com/p/375224982 由于前面兩篇文章已經(jīng)分析的非常詳細了，本文只簡要描述原理和具體說明用法。 考慮到訓(xùn)練過程中梯度幅值大部分是非常小的，故訓(xùn)練默認是 FP32 格式，如果能直接以 FP16 格式精度進行訓(xùn)練，理論上可以減少一半的內(nèi)存，達到加速訓(xùn)練和采用更大 batch size 的目的，但是直接以 FP16 訓(xùn)練會出現(xiàn)溢出問題，導(dǎo)致 NAN 或者參數(shù)更新失敗問題，而 AMP 的出現(xiàn)就是為了解決這個問題，其核心思想是 混合精度訓(xùn)練+動態(tài)損失放大：  1. 維護一個 FP32 數(shù)值精度模型的副本 2. 在每個 iteration a. 拷貝并且轉(zhuǎn)換成 FP16 模型 b. 前向傳播（FP16 的模型參數(shù))，此時 weights, activations 都是 FP16 c.loss 乘 scale factor s d. 反向傳播（FP16 的模型參數(shù)和參數(shù)梯度)， 此時 gradients 也是 FP16 e.參數(shù)梯度乘 1/s f.利用 FP16 的梯度更新 FP32 的模型參數(shù) 在 MMCV 中使用 AMP 分成兩種情況： - 在 OpenMMLab 上游庫例如 MMDetection 中使用 MMCV 的 AMP - 用戶只想簡單調(diào)用 MMCV 中的 AMP，而不依賴上游庫 OpenMMLab 上游庫 如何使用 MMCV 的 AMP 以 MMDectection 為例，用法非常簡單，只需要在配置中設(shè)置： fp16 = dict(loss_scale=512.) # 表示靜態(tài) scale # 表示動態(tài) scale fp16 = dict(loss_scale='dynamic') # 通過字典形式靈活開啟動態(tài) scale fp16 = dict(loss_scale=dict(init_scale=512.,mode='dynamic')) 三種不同設(shè)置在大部分模型上性能都非常接近，如果不想設(shè)置 loss_scale，則可以簡單的采用 loss_scale='dynamic' 調(diào)用 MMCV 中的 AMP 直接調(diào)用 MMCV 中的 AMP，這通常意味著用戶可能在其他庫或者自己寫的代碼庫中支持 AMP 功能。 需要特別強調(diào)的是 PyTorch 官方僅僅在 1.6 版本及其之后版本中開始支持 AMP，而 MMCV 中的 AMP 支持 1.3 及其之后版本。如果你想在 1.3 或者 1.5 中使用 AMP，那么使用 MMCV 是個非常不錯的選擇。 使用 MMCV 的 AMP 功能，只需要遵循以下幾個步驟即可： 1. 將 auto_fp16 裝飾器應(yīng)用到 model 的 forward 函數(shù)上； 2. 設(shè)置模型的 fp16_enabled 為 True 表示開啟 AMP 訓(xùn)練，否則不生效； 3. 如果開啟了 AMP，需要同時配置對應(yīng)的 FP16 優(yōu)化器配置 Fp16OptimizerHook； 4. 在訓(xùn)練的不同時刻，調(diào)用 Fp16OptimizerHook，如果你同時使用了 MMCV 中的 Runner 模塊，那么直接將第 3 步的參數(shù)輸入到 Runner 中即可； 5. (可選) 如果對應(yīng)某些 OP 希望強制運行在 FP32 上，則可以在對應(yīng)位置引入 force_fp32 裝飾器。
注意 force_fp32 要生效，依然需要 fp16_enabled 為 True 才生效。 2. 大 Batch 訓(xùn)練（梯度累加） 大 Batch 訓(xùn)練通常也稱為梯度累加策略，通常 PyTorch 一次迭代訓(xùn)練流程為： y_pred = model(xx)loss = loss_fn(y_pred, y)loss.backward()optimizer.step() optimizer.zero_grad()而梯度累加策略下常見的一次迭代訓(xùn)練流程為：
其核心思想就是對前幾次梯度進行累加，然后再統(tǒng)一進行參數(shù)更新，從而變相實現(xiàn)大 batch size 功能。需要注意的是如果模型中包括 BN 等考慮 batch 信息的層，那么性能可能會有輕微的差距。細節(jié)可以參考 https://github.com/open-mmlab/mmcv/pull/1221。 在 MMCV 中已經(jīng)實現(xiàn)了梯度累加功能，其核心代碼位于mmcv/runner/hooks/optimizer.py GradientCumulativeOptimizerHook 中，和 AMP 實現(xiàn)一樣是采用 Hook 實現(xiàn)的。使用方法和 AMP 類似，只需要將第一節(jié)中的 Fp16OptimizerHook 替換為 GradientCumulativeOptimizerHook 或者 GradientCumulativeFp16OptimizerHook 即可。 其核心實現(xiàn)如下所示： @HOOKS.register_module()class GradientCumulativeOptimizerHook(OptimizerHook): def __init__(self, cumulative_iters=1, **kwargs): self.cumulative_iters = cumulative_iters self.divisible_iters = 0 # 剩余的可以被 cumulative_iters 整除的訓(xùn)練迭代次數(shù) self.remainder_iters = 0 # 剩余累加次數(shù) self.initialized = False def after_train_iter(self, runner): # 只需要運行一次即可 if not self.initialized: self._init(runner) if runner.iter < self.divisible_iters: loss_factor = self.cumulative_iters else: loss_factor = self.remainder_iters loss = runner.outputs['loss'] loss = loss / loss_factor loss.backward() if (self.every_n_iters(runner, self.cumulative_iters) or self.is_last_iter(runner)): runner.optimizer.step() runner.optimizer.zero_grad() def _init(self, runner): residual_iters = runner.max_iters - runner.iter self.divisible_iters = ( residual_iters // self.cumulative_iters * self.cumulative_iters) self.remainder_iters = residual_iters - self.divisible_iters self.initialized = True 需要明白 divisible_iters 和 remainder_iters 的含義： 從頭訓(xùn)練 此時在開始訓(xùn)練時 iter=0，一共迭代 max_iters=102 次，梯度累加次數(shù)是 4，由于 102 無法被 4 整除，也就是最后的 102-(102 // 4)*4=2 個迭代是額外需要考慮的，在最后 2 個訓(xùn)練迭代中 loss_factor 不能除以 4，而是 2，這樣才是最合理的做法。其中 remainder_iters=2，divisible_iters=100，residual_iters=102。 resume 訓(xùn)練 假設(shè)在梯度累加的中途退出，然后進行 resume 訓(xùn)練，此時 iter 不是 0，由于優(yōu)化器對象需要重新初始化，為了保證剩余的不能被累加次數(shù)的訓(xùn)練迭代次數(shù)能夠正常計算，需要重新計算 residual_iters。 3. 梯度檢查點 梯度檢查點是一種用訓(xùn)練時間換取顯存的辦法，其核心原理是在反向傳播時重新計算神經(jīng)網(wǎng)絡(luò)的中間激活值而不用在前向時存儲，torch.utils.checkpoint 包中已經(jīng)實現(xiàn)了對應(yīng)功能。簡要實現(xiàn)過程是：在前向階段傳遞到 checkpoint 中的 forward 函數(shù)會以 torch.no_grad 模式運行，并且僅僅保存輸入?yún)?shù)和 forward 函數(shù)，在反向階段重新計算其 forward 輸出值。 具體用法非常簡單，以 ResNet 的 BasicBlock 為例：
self.with_cp 為 True，表示要開啟梯度檢查點功能。 checkpoint 在用法上面需要注意以下幾點： 1. 模型的第一層不能用 checkpoint 或者說 forward 輸入中不能所有輸入的 requires_grad 屬性都是 False，因為其內(nèi)部實現(xiàn)是依靠輸入的 requires_grad 屬性來判斷輸出返回是否需要梯度，而通常模型第一層輸入是 image tensor，其 requires_grad 通常是 False。一旦你第一層用了 checkpoint，那么意味著這個 forward 函數(shù)不會有任何梯度，也就是說不會進行任何參數(shù)更新，沒有任何使用的必要，具體見 https://discuss.pytorch.org/t/use-of-torch-utils-checkpoint-checkpoint-causes-simple-model-to-diverge/116271。如果第一層用了 checkpoint， PyTorch 會打印 None of the inputs have requires_grad=True. Gradients will be Non 警告； 2. 對于 dropout 這種 forward 存在隨機性的層，需要保證 preserve_rng_state 為 True (默認就是 True，所以不用擔(dān)心)，一旦標志位設(shè)置為 True，在 forward 會存儲 RNG 狀態(tài)，然后在反向傳播的時候讀取該 RNG，保證兩次 forward 輸出一致。如果你確定不需要保存 RNG，則可以設(shè)置 preserve_rng_state 為 False，省掉一些不必要的運行邏輯； 3. 其他注意事項，可以參考官方文檔 https://pytorch.org/docs/stable/checkpoint.html# 其核心實現(xiàn)如下所示： class CheckpointFunction(torch.autograd.Function): @staticmethod def forward(ctx, run_function, preserve_rng_state, *args): # 檢查輸入?yún)?shù)是否需要梯度 check_backward_validity(args) # 保存必要的狀態(tài) ctx.run_function = run_function ctx.save_for_backward(*args) with torch.no_grad(): # 以 no_grad 模型運行一遍 outputs = run_function(*args) return outputs @staticmethod def backward(ctx, *args): # 讀取輸入?yún)?shù) inputs = ctx.saved_tensors # Stash the surrounding rng state, and mimic the state that was # present at this time during forward. Restore the surrounding state # when we're done. rng_devices = [] with torch.random.fork_rng(devices=rng_devices, enabled=ctx.preserve_rng_state): # detach 掉當(dāng)前不需要考慮的節(jié)點 detached_inputs = detach_variable(inputs) # 重新運行一遍 with torch.enable_grad(): outputs = ctx.run_function(*detached_inputs) if isinstance(outputs, torch.Tensor): outputs = (outputs,) # 計算該子圖梯度 torch.autograd.backward(outputs, args) grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else inp for inp in detached_inputs) return (None, None) + grads4. 實驗驗證 為了驗證上述策略是否真的能夠省顯存，采用 mmdetection 庫進行驗證，基本環(huán)境如下：
(1) base - 數(shù)據(jù)集：pascal voc - 算法是 retinanet，對應(yīng)配置文件為 retinanet_r50_fpn_1x_voc0712.py - 為了防止 lr 過大導(dǎo)致訓(xùn)練出現(xiàn) nan，需要將 lr 設(shè)置為 0.01/8=0.00125 - bs 設(shè)置為 2 (2) 混合精度 AMP 在 base 配置基礎(chǔ)上新增如下配置即可： fp16 = dict(loss_scale=512.)(3) 梯度累加 在 base 配置基礎(chǔ)上替換 optimizer_config 為如下：
(4) 梯度檢查點 在 base 配置基礎(chǔ)上在 backbone 部分開啟 with_cp 標志即可： model = dict(backbone=dict(with_cp=True), bbox_head=dict(num_classes=20))每個實驗總共迭代 1300 次，統(tǒng)計占用顯存、訓(xùn)練總時長。  1. 對比 base 和 AMP 可以發(fā)現(xiàn)，由于實驗顯卡是不支持 AMP 的，故只能節(jié)省顯存，速度會特別慢，如果本身顯卡支持 AMP 則可以實現(xiàn)在節(jié)省顯存的同時提升訓(xùn)練速度； 2. 對比 base 和梯度累加可以發(fā)現(xiàn)，在相同 bs 情況下，梯度累加 2 次相當(dāng)于 bs 擴大一倍，但是顯存增加不多。如果將 bs 縮小一倍，則可以實現(xiàn)在相同 bs 情況下節(jié)省大概一倍顯存； 3. 對比 base 和梯度檢查點可以發(fā)現(xiàn)，可以節(jié)省一定的顯存，但是訓(xùn)練時長會增加一些。 從上面簡單實驗可以發(fā)現(xiàn)，AMP、梯度累加和梯度檢查點確實可以在不同程度減少顯存，而且這三個策略是正交的，可以同時使用。 本文簡要描述了三個在 MMCV 中集成且可以通過配置一行開啟的節(jié)省顯存策略，這三個策略比較常用也比較成熟。 隨著模型規(guī)模的不斷增長，也出現(xiàn)了很多新的策略，例如模型參數(shù)壓縮、動態(tài)顯存優(yōu)化、使用 CPU 內(nèi)存暫存策略以及分布式情況下 pytorch 1.10 最新支持的 ZeroRedundancyOptimizer 等等。 同時我們非常歡迎對計算機視覺前沿技術(shù)、開源項目開發(fā)有興趣的同學(xué)以全職或?qū)嵙?xí)的身份加入 OpenMMLab 團隊。歡迎大家聯(lián)系小助手投遞簡歷哦！  |
|
|
來自： LZS2851 > 《OpenMMMLab》
