拿什么拯救我的 4G 顯卡

隨著深度學(xué)習(xí)快速發(fā)展，同時(shí)伴隨著模型參數(shù)的爆炸式增長(zhǎng)，對(duì)顯卡的顯存容量提出了越來(lái)越高的要求，如何在單卡小容量顯卡上面訓(xùn)練模型是一直以來(lái)大家關(guān)心的問題。

本文結(jié)合 MMCV 開源庫(kù)對(duì)一些常用的節(jié)省顯存策略進(jìn)行簡(jiǎn)要分析，并希望能夠?qū)θ绾渭床寮从玫氖褂?MMCV 中相應(yīng)節(jié)省顯存策略的用戶提供一個(gè)簡(jiǎn)單的指引。

本文涉及到的 PyTorch 節(jié)省顯存的策略包括：

 - 混合精度訓(xùn)練

 - 大 batch 訓(xùn)練或者稱為梯度累加

 - gradient checkpointing 梯度檢查點(diǎn)

混合精度訓(xùn)練全稱為 Automatic Mixed Precision，簡(jiǎn)稱為 AMP，也就是我們常說(shuō)的 FP16。在前系列解讀中已經(jīng)詳細(xì)分析了 AMP 原理、源碼實(shí)現(xiàn)以及 MMCV 中如何一行代碼使用 AMP，具體鏈接見：

PyTorch 源碼解讀之 torch.cuda.amp: 自動(dòng)混合精度詳解：

https://zhuanlan.zhihu.com/p/348554267

OpenMMLab 中混合精度訓(xùn)練 AMP 的正確打開方式：

https://zhuanlan.zhihu.com/p/375224982

由于前面兩篇文章已經(jīng)分析的非常詳細(xì)了，本文只簡(jiǎn)要描述原理和具體說(shuō)明用法。

考慮到訓(xùn)練過程中梯度幅值大部分是非常小的，故訓(xùn)練默認(rèn)是 FP32 格式，如果能直接以 FP16 格式精度進(jìn)行訓(xùn)練，理論上可以減少一半的內(nèi)存，達(dá)到加速訓(xùn)練和采用更大 batch size 的目的，但是直接以 FP16 訓(xùn)練會(huì)出現(xiàn)溢出問題，導(dǎo)致 NAN 或者參數(shù)更新失敗問題，而 AMP 的出現(xiàn)就是為了解決這個(gè)問題，其核心思想是 混合精度訓(xùn)練+動(dòng)態(tài)損失放大：

1. 維護(hù)一個(gè) FP32 數(shù)值精度模型的副本

2. 在每個(gè) iteration                                                  

   a. 拷貝并且轉(zhuǎn)換成 FP16 模型

  b. 前向傳播（FP16 的模型參數(shù))，此時(shí) weights, activations 都是 FP16

   c.loss 乘 scale factor s

  d. 反向傳播（FP16 的模型參數(shù)和參數(shù)梯度)， 此時(shí) gradients 也是 FP16

   e.參數(shù)梯度乘 1/s

   f.利用 FP16 的梯度更新 FP32 的模型參數(shù)

在 MMCV 中使用 AMP 分成兩種情況：

 - 在 OpenMMLab 上游庫(kù)例如 MMDetection 中使用 MMCV 的 AMP

 - 用戶只想簡(jiǎn)單調(diào)用 MMCV 中的 AMP，而不依賴上游庫(kù)

以 MMDectection 為例，用法非常簡(jiǎn)單，只需要在配置中設(shè)置：

三種不同設(shè)置在大部分模型上性能都非常接近，如果不想設(shè)置 loss_scale，則可以簡(jiǎn)單的采用 loss_scale='dynamic'

直接調(diào)用 MMCV 中的 AMP，這通常意味著用戶可能在其他庫(kù)或者自己寫的代碼庫(kù)中支持 AMP 功能。

需要特別強(qiáng)調(diào)的是 PyTorch 官方僅僅在 1.6 版本及其之后版本中開始支持 AMP，而 MMCV 中的 AMP 支持 1.3 及其之后版本。如果你想在 1.3 或者 1.5 中使用 AMP，那么使用 MMCV 是個(gè)非常不錯(cuò)的選擇。

使用 MMCV 的 AMP 功能，只需要遵循以下幾個(gè)步驟即可：

1. 將 auto_fp16 裝飾器應(yīng)用到 model 的 forward 函數(shù)上；

2. 設(shè)置模型的 fp16_enabled 為 True 表示開啟 AMP 訓(xùn)練，否則不生效；

3. 如果開啟了 AMP，需要同時(shí)配置對(duì)應(yīng)的 FP16 優(yōu)化器配置 Fp16OptimizerHook；

4. 在訓(xùn)練的不同時(shí)刻，調(diào)用 Fp16OptimizerHook，如果你同時(shí)使用了 MMCV 中的 Runner 模塊，那么直接將第 3 步的參數(shù)輸入到 Runner 中即可；

5. (可選) 如果對(duì)應(yīng)某些 OP 希望強(qiáng)制運(yùn)行在 FP32 上，則可以在對(duì)應(yīng)位置引入 force_fp32 裝飾器。

# 1 作用到 forward 函數(shù)中class ExampleModule(nn.Module):
    @auto_fp16()    def forward(self, x, y):        return x, y        # 2 如果開啟 AMP，則需要加入開啟標(biāo)志model.fp16_enabled = True     
# 3 配置 Fp16OptimizerHookoptimizer_config = Fp16OptimizerHook(    **cfg.optimizer_config, **fp16_cfg, distributed=distributed)
# 4 傳遞給 runnerrunner.register_training_hooks(cfg.lr_config, optimizer_config,                               cfg.checkpoint_config, cfg.log_config,                               cfg.get('momentum_config', None))    # 5 可選class ExampleModule(nn.Module):
    @auto_fp16()    def forward(self, x, y):        features=self._forward(x, y)        loss=self._loss(features,labels)        return loss        def _forward(self, x, y):       pass            @force_fp32(apply_to=('features',))    def _loss(features,labels) :        pass                

注意 force_fp32 要生效，依然需要 fp16_enabled 為 True 才生效。

大 Batch 訓(xùn)練通常也稱為梯度累加策略，通常 PyTorch 一次迭代訓(xùn)練流程為：

而梯度累加策略下常見的一次迭代訓(xùn)練流程為：

y_pred = model(xx)loss = loss_fn(y_pred, y)
loss = loss / cumulative_itersloss.backward()
if current_iter % cumulative_iters==0    optimizer.step()     optimizer.zero_grad()

其核心思想就是對(duì)前幾次梯度進(jìn)行累加，然后再統(tǒng)一進(jìn)行參數(shù)更新，從而變相實(shí)現(xiàn)大 batch size 功能。需要注意的是如果模型中包括 BN 等考慮 batch 信息的層，那么性能可能會(huì)有輕微的差距。細(xì)節(jié)可以參考 https://github.com/open-mmlab/mmcv/pull/1221。

在 MMCV 中已經(jīng)實(shí)現(xiàn)了梯度累加功能，其核心代碼位于mmcv/runner/hooks/optimizer.py

GradientCumulativeOptimizerHook 中，和 AMP 實(shí)現(xiàn)一樣是采用 Hook 實(shí)現(xiàn)的。使用方法和 AMP 類似，只需要將第一節(jié)中的 Fp16OptimizerHook 替換為 GradientCumulativeOptimizerHook 或者 GradientCumulativeFp16OptimizerHook 即可。

其核心實(shí)現(xiàn)如下所示：

需要明白 divisible_iters 和 remainder_iters 的含義：

此時(shí)在開始訓(xùn)練時(shí) iter=0，一共迭代 max_iters=102 次，梯度累加次數(shù)是 4，由于 102 無(wú)法被 4 整除，也就是最后的 102-(102 // 4)*4=2 個(gè)迭代是額外需要考慮的，在最后 2 個(gè)訓(xùn)練迭代中 loss_factor 不能除以 4，而是 2，這樣才是最合理的做法。其中 remainder_iters=2，divisible_iters=100，residual_iters=102。

假設(shè)在梯度累加的中途退出，然后進(jìn)行 resume 訓(xùn)練，此時(shí) iter 不是 0，由于優(yōu)化器對(duì)象需要重新初始化，為了保證剩余的不能被累加次數(shù)的訓(xùn)練迭代次數(shù)能夠正常計(jì)算，需要重新計(jì)算 residual_iters。

梯度檢查點(diǎn)是一種用訓(xùn)練時(shí)間換取顯存的辦法，其核心原理是在反向傳播時(shí)重新計(jì)算神經(jīng)網(wǎng)絡(luò)的中間激活值而不用在前向時(shí)存儲(chǔ)，torch.utils.checkpoint 包中已經(jīng)實(shí)現(xiàn)了對(duì)應(yīng)功能。簡(jiǎn)要實(shí)現(xiàn)過程是：在前向階段傳遞到 checkpoint 中的 forward 函數(shù)會(huì)以 torch.no_grad 模式運(yùn)行，并且僅僅保存輸入?yún)?shù)和 forward 函數(shù)，在反向階段重新計(jì)算其 forward 輸出值。

具體用法非常簡(jiǎn)單，以 ResNet 的 BasicBlock 為例：

def forward(self, x):    def _inner_forward(x):        identity = x        out = self.conv1(x)        out = self.norm1(out)        out = self.relu(out)        out = self.conv2(out)        out = self.norm2(out)        if self.downsample is not None:            identity = self.downsample(x)        out += identity        return out            # x.requires_grad 這個(gè)判斷很有必要    if self.with_cp and x.requires_grad:        out = cp.checkpoint(_inner_forward, x)    else:        out = _inner_forward(x)    out = self.relu(out)    return out

self.with_cp 為 True，表示要開啟梯度檢查點(diǎn)功能。

checkpoint 在用法上面需要注意以下幾點(diǎn)：

1. 模型的第一層不能用 checkpoint 或者說(shuō) forward 輸入中不能所有輸入的 requires_grad 屬性都是 False，因?yàn)槠鋬?nèi)部實(shí)現(xiàn)是依靠輸入的 requires_grad 屬性來(lái)判斷輸出返回是否需要梯度，而通常模型第一層輸入是 image tensor，其 requires_grad 通常是 False。一旦你第一層用了 checkpoint，那么意味著這個(gè) forward 函數(shù)不會(huì)有任何梯度，也就是說(shuō)不會(huì)進(jìn)行任何參數(shù)更新，沒有任何使用的必要，具體見 https://discuss.pytorch.org/t/use-of-torch-utils-checkpoint-checkpoint-causes-simple-model-to-diverge/116271。如果第一層用了 checkpoint， PyTorch 會(huì)打印 None of the inputs have requires_grad=True. Gradients will be Non 警告；

2. 對(duì)于 dropout 這種 forward 存在隨機(jī)性的層，需要保證 preserve_rng_state 為 True (默認(rèn)就是 True，所以不用擔(dān)心)，一旦標(biāo)志位設(shè)置為 True，在 forward 會(huì)存儲(chǔ) RNG 狀態(tài)，然后在反向傳播的時(shí)候讀取該 RNG，保證兩次 forward 輸出一致。如果你確定不需要保存 RNG，則可以設(shè)置 preserve_rng_state 為 False，省掉一些不必要的運(yùn)行邏輯；

3. 其他注意事項(xiàng)，可以參考官方文檔

 https://pytorch.org/docs/stable/checkpoint.html#

其核心實(shí)現(xiàn)如下所示：

為了驗(yàn)證上述策略是否真的能夠省顯存，采用 mmdetection 庫(kù)進(jìn)行驗(yàn)證，基本環(huán)境如下：

顯卡: GeForce GTX 1660PyTorch: 1.7.1CUDA Runtime 10.1MMCV: 1.3.16MMDetection: 2.17.0

(1) base

- 數(shù)據(jù)集：pascal voc

- 算法是 retinanet，對(duì)應(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ǔ)上新增如下配置即可：

(3) 梯度累加

在 base 配置基礎(chǔ)上替換 optimizer_config 為如下：

# 累加2次optimizer_config = dict(type='GradientCumulativeOptimizerHook', cumulative_iters=2)

(4) 梯度檢查點(diǎn)

在 base 配置基礎(chǔ)上在 backbone 部分開啟 with_cp 標(biāo)志即可：

每個(gè)實(shí)驗(yàn)總共迭代 1300 次，統(tǒng)計(jì)占用顯存、訓(xùn)練總時(shí)長(zhǎng)。

1. 對(duì)比 base 和 AMP 可以發(fā)現(xiàn)，由于實(shí)驗(yàn)顯卡是不支持 AMP 的，故只能節(jié)省顯存，速度會(huì)特別慢，如果本身顯卡支持 AMP 則可以實(shí)現(xiàn)在節(jié)省顯存的同時(shí)提升訓(xùn)練速度；

2. 對(duì)比 base 和梯度累加可以發(fā)現(xiàn)，在相同 bs 情況下，梯度累加 2 次相當(dāng)于 bs 擴(kuò)大一倍，但是顯存增加不多。如果將 bs 縮小一倍，則可以實(shí)現(xiàn)在相同 bs 情況下節(jié)省大概一倍顯存；

3. 對(duì)比 base 和梯度檢查點(diǎn)可以發(fā)現(xiàn)，可以節(jié)省一定的顯存，但是訓(xùn)練時(shí)長(zhǎng)會(huì)增加一些。

從上面簡(jiǎn)單實(shí)驗(yàn)可以發(fā)現(xiàn)，AMP、梯度累加和梯度檢查點(diǎn)確實(shí)可以在不同程度減少顯存，而且這三個(gè)策略是正交的，可以同時(shí)使用。