隨著深度學(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) + grads 4. 實驗驗證 為了驗證上述策略是否真的能夠省顯存,采用 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》