视频训练效率太低？试试这个加速算法

我们都知道，在视频上训练深度网络 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 训练加速算法的具体细节。

论文链接：

https://openaccess.thecvf.com/content_CVPR_2020/papers/Wu_A_Multigrid_Method_for_Efficiently_Training_Video_Models_CVPR_2020_paper.pdf

算法介绍✦

在 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 。Long cycle 与 step-wise 学习率衰减策略同步，即对于每个学习率阶段，这四种 shape 均经历一次， 而且对每个 shape 进行相同 iteration 次数的迭代训练。

在训练的 Iteraions 中，每段固定的 Iterations 分别进行时序的采样，采样的同时限制长宽比，四段采样&长宽即为图（b）中的四段，从 x x x 到 x x x 。

Short cycle

Short cycle 在 每一个 iteration 都会改变输入的空间形状 H 和 W ， B 也会随之改变。

如果 mod 3 = 0, H 和 W 设置为 x , B 设为 。

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

如果 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上。

实验结果✦

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

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

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

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 ！