拿什么拯救我的 4G 顯卡： PyTorch 節(jié)省顯存的策略總結(jié)

0 前言

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

• 混合精度訓(xùn)練
• 大 batch 訓(xùn)練或者稱(chēng)為梯度累加
• gradient checkpointing 梯度檢查點(diǎn)

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

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

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

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

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

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

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

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

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

• 拷貝并且轉(zhuǎn)換成 FP16 模型
• 前向傳播（FP16 的模型參數(shù))，此時(shí) weights, activations 都是 FP16
• loss 乘 scale factor s
• 反向傳播（FP16 的模型參數(shù)和參數(shù)梯度)， 此時(shí) gradients 也是 FP16
• 參數(shù)梯度乘 1/s
• 利用 FP16 的梯度更新 FP32 的模型參數(shù)

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

• 在 OpenMMLab 上游庫(kù)例如 MMDetection 中使用 MMCV 的 AMP
• 用戶(hù)只想簡(jiǎn)單調(diào)用 MMCV 中的 AMP，而不依賴(lài)上游庫(kù)

(1) OpenMMLab 上游庫(kù)如何使用 MMCV 的 AMP

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

fp16 = dict(loss_scale=512.) # 表示靜態(tài) scale 

# 表示動(dòng)態(tài) scale 
fp16 = dict(loss_scale='dynamic')  

# 通過(guò)字典形式靈活開(kāi)啟動(dòng)態(tài) scale 
fp16 = dict(loss_scale=dict(init_scale=512.,mode='dynamic'))  

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

(2) 調(diào)用 MMCV 中的 AMP

直接調(diào)用 MMCV 中的 AMP，這通常意味著用戶(hù)可能在其他庫(kù)或者自己寫(xiě)的代碼庫(kù)中支持 AMP 功能。需要特別強(qiáng)調(diào)的是 PyTorch 官方僅僅在 1.6 版本及其之后版本中開(kāi)始支持 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 表示開(kāi)啟 AMP 訓(xùn)練，否則不生效
3. 如果開(kāi)啟了 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 如果開(kāi)啟 AMP，則需要加入開(kāi)啟標(biāo)志
model.fp16_enabled = True     

# 3 配置 Fp16OptimizerHook
optimizer_config = Fp16OptimizerHook(
    **cfg.optimizer_config, **fp16_cfg, distributed=distributed)

# 4 傳遞給 runner
runner.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 才生效。

2 大 Batch 訓(xùn)練（梯度累加）

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

y_pred = model(xx)
loss = loss_fn(y_pred, y)
loss.backward()
optimizer.step() 
optimizer.zero_grad()

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

y_pred = model(xx)
loss = loss_fn(y_pred, y)

loss = loss / cumulative_iters
loss.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 類(lèi)似，只需要將第一節(jié)中的 Fp16OptimizerHook 替換為 GradientCumulativeOptimizerHook 或者 GradientCumulativeFp16OptimizerHook 即可。其核心實(shí)現(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):
        # 只需要運(yùn)行一次即可
        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 的含義：

(1) 從頭訓(xùn)練

此時(shí)在開(kāi)始訓(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。

(2) resume 訓(xùn)練

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

3 梯度檢查點(diǎn)

梯度檢查點(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)過(guò)程是：在前向階段傳遞到 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，表示要開(kāi)啟梯度檢查點(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ù)更新，沒(méi)有任何使用的必要，具體見(jiàn) 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)如下所示：

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 模型運(yùn)行一遍
            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é)點(diǎn)
            detached_inputs = detach_variable(inputs)
            # 重新運(yùn)行一遍
            with torch.enable_grad():
                outputs = ctx.run_function(*detached_inputs)
       
        if isinstance(outputs, torch.Tensor):
            outputs = (outputs,)
        # 計(jì)算該子圖梯度
        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 實(shí)驗(yàn)驗(yàn)證

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

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

(1) base

• 數(shù)據(jù)集：pascal voc
• 算法是 retinanet，對(duì)應(yīng)配置文件為 retinanet_r50_fpn_1x_voc0712.py
• 為了防止 lr 過(guò)大導(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 為如下：

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

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

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

model = dict(backbone=dict(with_cp=True),
             bbox_head=dict(num_classes=20))

每個(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í)使用。

5 總結(jié)

本文簡(jiǎn)要描述了三個(gè)在 MMCV 中集成且可以通過(guò)配置一行開(kāi)啟的節(jié)省顯存策略，這三個(gè)策略比較常用也比較成熟。隨著模型規(guī)模的不斷增長(zhǎng)，也出現(xiàn)了很多新的策略，例如模型參數(shù)壓縮、動(dòng)態(tài)顯存優(yōu)化、使用 CPU 內(nèi)存暫存策略以及分布式情況下 PyTorch 1.10 最新支持的 ZeroRedundancyOptimizer 等等。

快速鏈接直達(dá) MMCV 算法庫(kù)，歡迎大家 Star：

https://github.com/open-mmlab/mmcv