视频训练效率太低？Multigrid 加速算法了解一下

0.前言

我们都知道，在视频上训练深度网络 3D CNN 比训练 2D CNN 图像模型的计算量更大，可能要大一个数量级。 长时间的训练会消耗大量的硬件和资源，在减缓视频理解研究领域发展的同时，也会阻碍该领域在真实场景的应用。
在 3D CNN 视频模型中， 每一个 mini-batch 的输入 shape 为 Batch_size x T (采样帧数) x H (高度） x W（宽度）， 通常在训练中 Batch_size，T，H，W 的值都是固定的。
为了解决训练效率的问题，论文 "A Multigrid Method for Efficiently Training Video Models" 提出了一种动态改变 Batch_size、采样帧数 T、 每一帧的宽度 W 和高度 H 的方式，能在保证训练精度不变的情况下，加速训练收敛。
今天我们将解读该篇论文，带领大家一起了解 Multigrid 训练加速算法的具体细节。

1.算法介绍

在 3D CNN 视频模型中， 每一个 mini-batch 的输入 shape 为 Batch_size x T (采样帧数) x H (高度） x W（宽度）， 通常在训练中 Batch_size，T，H，W 的值都是固定的。增大 T，H，W 维度通常可以提高视频模型的 accuracy， 而增加 Batch_size 且 减小 T, H, W 通常可以换来更快的训练但是 accuracy 会降低。受到数值计算领域中通过改变 coarse and fine grids 来解决优化问题思想的启发，Multigrid 算法在保证 B x T x H x W 的值固定的前提下， 动态改变 T x H x W 和 B 的值， 在训练早期使用具有相对较小 T, H, W（“粗网格”）的大 mini-batch，后期使用具有较大 T, H, W（“细网格”）的小 mini-batch，那么 SGD 可能能够更快地扫描数据，在保证精度的同时减少了训练时间。文章主要提出了 3 种网格策略： long cycle， short cycle 和 long + short cycle 。

Long cycle

采用 4 种 mini-batch shapes:  x   x   x   ，  x   x   x  ，  x  x  x ，  x  x  x  。 Long cycle 与 step-wise 学习率衰减策略同步，即对于每个学习率阶段，这四种 shape 均经历一次， 而且对每个 shape 进行相同 iteration 次数的迭代训练。
在训练的 Iteraions 中，每段固定的Iterations 分别进行时序的采样，采样的同时限制长宽比，四段采样&长宽即为图（b）中的四段，从  x  x  x  ，  x  x  x  。

Short cycle

Short cycle 在 每一个 iteration  都会改变输入的 空间形状 H 和 W ， B 也会随之改变 。
1.如果  mod 3 = 0, H 和 W 设置为  x  , B 设为 

2.如果  mod 3 = 1, H 和 W 设置为  x  , B 设为 

3.如果  mod 3 = 2, H 和 W 设置为  x  ， B 设为 

Long + short cycle

将 long cycle 和 short cycle 策略结合， 对于每个学习率阶段，long cycle 的四种 shape 均经历一次且进行相同 iteration 次数的迭代训练， 同时在每种 shape 的 iterations 中使用 short cycle， 交替改变 B, H 和 W 的值。

学习率放大

当 mini-batch 大小由于 long cycle 策略而发生变化时，我们根据 mini-batch 大小缩放因子（8×、4×、2× 或 1× ）来调整相应学习率的大小。我们仅在 long cycle 策略 shape 发生变化时调整学习率。

Fine-tune 阶段

如果 baseline 优化器使用 L 个学习率（LR）阶段，那么我们在前 L - 1 个 LR 阶段应用 long cycle 和 short cycle 策略。 我们使用相应的第 L 个阶段进行微调，以帮助匹配训练和测试分布。 在微调迭代的前半部分，我们使用 L-1 个阶段的 学习率，在后半部分，我们使用最终第 L 个阶段的学习率。 在微调时，我们仅使用 short cycle 作为数据增强。

Batch Normalization

Mini-batch size 是一个会影响 Batch Normalization 的超参数。由于 multigrid 方法使用动态变化的 mini-batch size，因此希望将其对 Batch Normalization 的影响与训练加速的影响分离。所以在实际使用中， 我们使用 Sub Batch Normalization 替代 Batch Normalization。 Sub Batch Normalization 会把 batch 维度分为 N 个 splits, 然后在每一个 split 里分别运行 Batch Normalization, 在每个样本的子集（batch 的 1/N）中单独计算统计信息。在 eval 的时候， Sub Batch Normalization 会把所有 splits 的统计信息汇集到一个 Batch Noramlization上。

2.实验结果

使用 R50-SlowFast 在 Kinetics-400 数据集上的精度和训练速度比较如下：

可以看到， 使用 Multigrid 训练策略可以在保持模型精度不变的同时加速 4-5 倍左右。

目前 MMAction2 v0.22.0 已经支持了 Multigrid 训练加速策略，接下来我们将带领大家一起学习 Multigrid 在MMAction2 的使用方法。

3.MMAction2 使用方法

MMAction2 目前支持了 SlowFast 模型在 Kinetics400 数据集上的 Multigrid 训练加速策略。（configs/recognition/slowfast/slowfast_multigrid_r50_8x8x1_358e_kinetics400_rgb.py)

使用步骤如下：

1.使用 step lr 学习率更新策略

lr_config = dict(policy='step', step=[94, 154, 196]) 

2.加入 multigrid 策略的设置

multigrid = dict( 
    long_cycle=True, 
    short_cycle=True, 
    epoch_factor=1.5, 
    long_cycle_factors=[[0.25, 0.7071],[0.5, 0.7071],[0.5, 1],[1, 1]], 
    short_cycle_factors=[0.5, 0.7071], 
    default_s=(224, 224), 
) 

3.使用 Sub Batch Normalization 替代 Batch Normalization，实验中发现使用或者不使用 Sub Batch Normalization 结果差别不大

model = dict( 
    type='Recognizer3D', 
    backbone=dict( 
        type='ResNet3dSlowFast', 
        pretrained=None, 
        resample_rate=4,  # tau 
        speed_ratio=4,  # alpha 
        channel_ratio=8,  # beta_inv 
        slow_pathway=dict( 
            type='resnet3d', 
            depth=50, 
            pretrained=None, 
            lateral=True, 
            fusion_kernel=7, 
            conv1_kernel=(1, 7, 7), 
            dilations=(1, 1, 1, 1), 
            conv1_stride_t=1, 
            pool1_stride_t=1, 
            inflate=(0, 0, 1, 1), 
            norm_cfg=dict(type='SubBatchNorm3D'), 
            norm_eval=False), 
        fast_pathway=dict( 
            type='resnet3d', 
            depth=50, 
            pretrained=None, 
            lateral=False, 
            base_channels=8, 
            conv1_kernel=(5, 7, 7), 
            conv1_stride_t=1, 
            pool1_stride_t=1, 
            norm_cfg=dict(type='SubBatchNorm3D'), 
            norm_eval=False)), 
    cls_head=dict( 
        type='SlowFastHead', 
        in_channels=2304,  # 2048+256 
        num_classes=400, 
        spatial_type='avg', 
        dropout_ratio=0.5), 
    # model training and testing settings 
    train_cfg=None, 
    test_cfg=dict(average_clips='prob')) 

4.使用 precise bn

precise_bn = dict(num_iters=200, interval=3) 

以下是使用 16 个GPU, 在 SlowFast 模型 和 Kinetics400 数据集上使用 Multigrid 训练策略的训练命令：

GPUS=16 tools/slurm_train.sh partition slowfast_multigrid_k400 configs/recognition/slowfast/slowfast_multigrid_r50_8x8x1_358e_kinetics400_rgb.py --work-dir work_dirs/slowfast_multigrid_r50_8x8x1_358e_kinetics400_rgb 

本期的 Multigrid 算法解读就到这里啦~你学会了吗？如果大家想要进行更深入地学习，可以借助我们 MMAction2 的算法库，非常欢迎大家来使用，Star, Fork 和 PR ！

https://github.com/open-mmlab/mmaction2​github.com/open-mmlab/mmaction2