显著提升模型精度!一文教你加载 MMRazor 轻量级骨干网络应用下游任务--以 MMYOLO 为例

673 阅读9分钟

MMYOLO 是一个基于 PyTorch 和 MMDetection 的 YOLO 系列工业核心算法库,为用户提供统一全面的评测流程、轻松可定制的模块组件以及支持多任务且高效的训练部署流程。

MMRazor 是 OpenMMLab 生态中面向模型压缩的开源算法库,目前主要涵盖了网络结构搜索、知识蒸馏、剪枝三类算法。

在本文中,今天我们将介绍如何在 OpenMMLab下游任务框架 MMYOLO 中调用从 MMRazor 中搜索出来的子网络作为 backbone。

NAS 前言

神经网络结构设计(Neural Architecture Design)是深度学习领域中最重要的研究方向之一。针对特定任务或场景的定制化深度网络通常需要研究员进行繁琐的针对性设计,耗时较长。在工业界,模型上线后需通过数据回流进行网络的迭代更新,而人工设计网络会极大拖延交付周期,影响算法在应用场景的最终收益。

NAS 的设计初衷为:高效地搜索,得到在参数量、计算量或推理时延等方面表现优秀的神经网络,因而 NAS 需兼顾搜索效率与适用范围。NAS 通常会在预定义的搜索空间中,通过各类优化算法进行自动搜索,因而 NAS 技术的研究重点为搜索空间的设计与优化范式的迭代。

有鉴于此,研究人员进一步提出“权重共享(Weight-sharing)”的设计思路,构建并训练一个超网络,而算法采样所得的子网络可直接继承超网络训练后的权重而无需从头开始训练,性能评估阶段所用时间大幅减少,搜索效率因而大幅提高。在后续的发展中,权重共享成为了 NAS 的基础设计。

在权重共享的基础上,One-Shot NAS 算法中的超网络被进一步设计为包含搜索空间中所有候选子网络,其本身表征了搜索空间。超网络本身的构造方式有很多种,如基于线性结构、基于有向无环图(DAG)等。超网络包含了架构搜索需要的模块与连接方式等信息,而其中的子网络数量极多,为简化搜索,通常需结合任务相关的先验知识,以减小搜索规模。然而,这会引入研究员的偏见,且超网络的设计对于缺乏领域知识的研究员而言是个挑战。

下游任务实战教程

MMRazor 近期贡献了一系列网络结构搜索算法的复现,例如 SPOS、Bignas、Once-For-All 等,在约束了计算量/参数量的搜索空间下,可得到精度优异的轻量级骨干网络,过去一年,无论是 YOLOV6 中的 EfficientRep,还是 PP-PicoDet 中 的 ESNet,都无一例外地通过 NAS 的思想重构基础骨干网络的设计,从而实现轻量级目标检测器性能的提升。

MMRazor 提供了基础的 backbone 用于构建常见的超网络,其大多由主流的 backbone 网络拓展而来,例如 AttentiveMobileNetV3、 SearchableShuffleNetV2 等。如下所示,在配置文件中通过指定 yaml 的模型配置文件进而可以加载超网络中的任意子网。

fix_subnet = 'https://download.openmmlab.com/mmrazor/v1/spos/spos_shufflenetv2_subnet_8xb128_in1k_flops_0.33M_acc_73.87_20211222-1f0a0b4d_subnet_cfg_v3.yaml'
nas_backbone = dict(
    type='mmrazor.sub_model',
    fix_subnet=fix_subnet,
    cfg = dict(
        _scope_='mmrazor',
        type='SearchableShuffleNetV2',
        widen_factor=1.0,
        arch_setting=_base_.arch_setting),  # 搜索空间对应配置
    extra_prefix='architecture.backbone.')  # 可以直接build出下游repo可用的backbone实例

基于 Weight-sharing 机制的 One-Shot NAS 应用于下游任务的流程如下表所示:

框架使用MMRazorMMRazorMMYOLO
主要输入输出超网络训练 (supernet.pth)子网搜索(subnet.pth/ subnet.yaml)下游子网权重加载(load subnet.pth & subnet.yaml)

目前,基于 Weight-sharing 机制的 One-Shot NAS 逐渐成为 NAS 主流设计范式。在实际生产环境中,上述流程方法可用于下游任务中,使其展现出更多的性能优势,以 bignas / spos / once-for-all 等 One-Shot NAS 算法搜索出的子网为例,给出了(通过在 Imagenet 上训练得到的 backbone 结合NAS 算法搜索到的子网) 应用到下游任务 MMYOLO 的配置文件以及训练结果。

接下来,我们就以 MMYOLO 为例,介绍 MMRazor 轻量级骨干网络应用下游任务实战教程。

MMRazor 可搜索子网配置

从 NAS 相关算法的配置文件中,我们可以观察到除了 .py 的配置文件之外,针对具体网络结构的配置文件均以 .yaml 格式保存在各个算法文件夹中,每一个 yaml 配合 subnet.py 将是一个独立的子网结构。
目前 MMRazor 主要支持的算法以及提供的 Backbone 如下:

注: 通过继承 SearchableBackbone 基类的可搜索网络结构可以配置包含不同超参数(层数、通道、卷积核大小等)的子网。

# mmrazor/configs/nas/mmcls/
├── bignas
│   ├── attentive_mobilenet_subnet_8xb256_in1k.py
│   ├── ATTENTIVE_SUBNET_A0.yaml
│   ├── ATTENTIVE_SUBNET_A6.yaml
│   └── README.md
├── onceforall
│   ├── ofa_mobilenet_subnet_8xb256_in1k.py
│   ├── OFA_SUBNET_NOTE8_LAT22.yaml
│   ├── OFA_SUBNET_NOTE8_LAT31.yaml
│   └── README.md
└── spos
    ├── spos_shufflenet_subnet_8xb128_in1k.py
    └── SPOS_SUBNET.yaml

以 ATTENTIVE_SUBNET_A0.yaml 以及 SPOS_SUBNET.yaml 为例,由于可搜索的 mutable 类型不同(bignas 算法中是 mutablevalue 和mutablechannel ,spos 算法中是 mutablemodule),在 yaml 文件中, key 为 mutable 的 name,而 value 是其对应搜索空间的取值。

以下举例说明:

  • 在 ATTENTIVE_SUBNET_A0.yaml 中,可搜索的通道数最终 backbone.first_channels 在该层通道选择 16。
# ATTENTIVE_SUBNET_A0.yaml
backbone.first_channels:
  chosen: 16
backbone.last_channels:
  chosen: 1792
backbone.layers.1.kernel_size:
  chosen: 3
backbone.layers.1.expand_ratio:
  chosen: 1
backbone.layers.1.depth:
  chosen: 1
backbone.layers.1.out_channels:
  chosen: 16
input_shape:
  chosen:
  - 192
  - 192
  • 在 SPOS_SUBNET.yaml 中,backbone.layers.0.0 的 candidates 最终选择 shuffle_7x7。
# SPOS_SUBNET.yaml 
backbone.layers.0.0:
  chosen: shuffle_7x7
backbone.layers.0.1:
  chosen: shuffle_3x3
backbone.layers.0.2:
  chosen: shuffle_7x7
backbone.layers.0.3:
  chosen: shuffle_3x3
backbone.layers.1.0:
  chosen: shuffle_xception
backbone.layers.1.1:
  chosen: shuffle_5x5
backbone.layers.1.2:
  chosen: shuffle_5x5
backbone.layers.1.3:
  chosen: shuffle_3x3

MMYOLO 下游任务训练配置详解

Yolov5 算法在社区有广泛的应用场景,首先我们先以 YOLOv5-s 直接替换为 SPOS 子网搜索的 SearchableShuffleNetV2 为例,通过对 backbone 以及 neck/head 的计算量的统计,我们发现参数量和计算量均低于原有网络配置。

通过对 nas_backbone 的 out_indices 输出的调整(分别为下采样 8、16、32 倍),其对应的通道改为 [160, 320, 640],由于通道需要与 neck/head 对齐,widen_factor 调整为 1.0:

_base_ = [
    'mmrazor::_base_/nas_backbones/spos_shufflenet_supernet.py',
    '../yolov5/yolov5_s-v61_syncbn_fast_8xb16-300e_coco.py'
]

checkpoint_file = 'https://download.openmmlab.com/mmrazor/v1/spos/spos_shufflenetv2_subnet_8xb128_in1k_flops_0.33M_acc_73.87_20211222-1f0a0b4d_v3.pth'  # noqa
fix_subnet = 'https://download.openmmlab.com/mmrazor/v1/spos/spos_shufflenetv2_subnet_8xb128_in1k_flops_0.33M_acc_73.87_20211222-1f0a0b4d_subnet_cfg_v3.yaml'  # noqa

widen_factor = 1.0
channels = [160, 320, 640] # 针对shufflenet的输出通道对齐到Neck

_base_.nas_backbone.out_indices = (1, 2, 3)
_base_.nas_backbone.init_cfg = dict(
    type='Pretrained',
    checkpoint=checkpoint_file,
    prefix='architecture.backbone.')
nas_backbone = dict(
    type='mmrazor.sub_model',
    fix_subnet=fix_subnet,
    cfg=_base_.nas_backbone,
    extra_prefix='architecture.backbone.')

_base_.model.backbone = nas_backbone
_base_.model.neck.widen_factor = widen_factor
_base_.model.neck.in_channels = channels # 注意:SearchableShuffleNetV2 输出的3个通道是 [160, 320, 640],和原先的 yolov5-s neck 不匹配,需要更改
_base_.model.neck.out_channels = channels
_base_.model.bbox_head.head_module.in_channels = channels  # head 部分输入通道也要做相应更改
_base_.model.bbox_head.head_module.widen_factor = widen_factor

find_unused_parameters = True

此时如下表所示,在 backbone 上参数量更少,部分参数量由于通道的扩增填补在 neck/head 中,此时整网的参数量可以控制在与原模型相同的规模,在无需更改其余训练配置项的前提下,在 coco 上 mAP 提升 0.2。

Table(Params; Flops)BackboneNeck + headmodelCoco mAP
Baseline(yolov5_s-v61)4.171M;5.212G3.064M;3.053G7.235M; 8.265G37.7
+ spos_shufflenetv2 Backbone2.329M; 2.378G4.713M;4.562G7.042M; 7.03G37.9 (+ 0.2)

选择 NAS 可搜索的子网 backbone 时,同时要将 nas_backbone 的类型指定为 mmrazor.sub_model, 此时 fix_subnet 必须指定,yaml 可以是 MMRazor 仓库下提供的默认配置,也可以是根据实际需求重新搜索得到的 yaml,与此同时,也要注意相应的 checkpoint_file 也需要和 yaml 匹配上。

 MMRazor 中选择适合的子网模型结构

在使用 NAS 搜索的权重和配置不满足需求的时候,我们需要重新在 MMRazor 中 重启 SearchLoop 搜索流程,利用不同算法的搜索阶段随机采样 100 个模型如下,我们可以观察不同模型大小的范围如下:

其中 spos/ofa 子网需要 retrain 或者 finetune 步骤,因此直接超网搜索出子网精度较差。

不同搜索的空间的计算量&参数量如下图所示:

以 YOLOv6-s 为例,YOLOv6EfficientRep backbone 参数量大于 10M,此时只有 AttentiveNAS 中的搜索空间可以提供参数量大于 10M 的 backbone,选用 Attentive_A6 的配置文件。如果可替换的 backbone 参数量在其他范围内,可以再选用合适的搜索空间,然后进行针对性的限制条件进行搜索。以 Bignas 系列算法为例,超网无需 retrain 即可得到子网权重,那么直接可以在限制 constraints_range 字段中 Params/Flops的参数,实现了子网在 Imagenet 上对验证集 Top-1 准确率的搜索。搜索完成后得到 best_fix_subnet.yaml 以及相应的子网权重 subnet_2023XXXX_0725.pth

# https://github.com/open-mmlab/mmrazor/blob/dev-1.x/configs/nas/mmcls/bignas/attentive_mobilenet_search_8xb128_in1k.py
_base_ = ['./attentive_mobilenet_supernet_32xb64_in1k.py']

train_cfg = dict(
    _delete_=True,
    type='mmrazor.EvolutionSearchLoop',
    dataloader=_base_.val_dataloader,
    evaluator=_base_.val_evaluator,
    max_epochs=20,
    num_candidates=50,
    top_k=10,
    num_mutation=25,
    num_crossover=25,
    mutate_prob=0.1,
    calibrate_sample_num=4096,
    constraints_range=dict(params=(14., 16.)), # 针对性地对模型大小进行限制
    score_key='accuracy/top1')

通过对模型配置文件 best_fix_subnet.yaml 以及相应的权重 subnet_2023XXXX_0725.pth 进行加载,首先可以根据下采样 8、16、32 倍确定输出通道数。打印 AttentiveMobileNetV3 每层输出通道为 [40, 128, 224]。在衔接的 neck / head 部分对齐相应的通道输出。

_base_ = [
    'mmrazor::_base_/nas_backbones/attentive_mobilenetv3_supernet.py',
    '../yolov6/yolov6_l_syncbn_fast_8xb32-300e_coco.py'
]

checkpoint_file = 'https://download.openmmlab.com/mmrazor/v1/bignas/attentive_mobilenet_subnet_8xb256_in1k_flops-0.93G_acc-80.81_20221229_200440-73d92cc6.pth'  # noqa
# 加载搜索所得yaml配置
fix_subnet = 'https://download.openmmlab.com/mmrazor/v1/bignas/ATTENTIVE_SUBNET_A6.yaml'  # noqa
deepen_factor = 1.2
widen_factor = 1
channels = [40, 128, 224]
mid_channels = [40, 128, 224]

_base_.train_dataloader.batch_size = 16
_base_.nas_backbone.out_indices = (2, 4, 6)
_base_.nas_backbone.conv_cfg = dict(type='mmrazor.BigNasConv2d')
_base_.nas_backbone.norm_cfg = dict(type='mmrazor.DynamicBatchNorm2d')
_base_.nas_backbone.init_cfg = dict(
    type='Pretrained',
    checkpoint=checkpoint_file,
    prefix='architecture.backbone.')
nas_backbone = dict(
    type='mmrazor.sub_model',
    fix_subnet=fix_subnet,
    cfg=_base_.nas_backbone,
    extra_prefix='backbone.')

_base_.model.backbone = nas_backbone
_base_.model.neck.widen_factor = widen_factor
_base_.model.neck.deepen_factor = deepen_factor
_base_.model.neck.in_channels = channels
_base_.model.neck.out_channels = mid_channels
_base_.model.bbox_head.head_module.in_channels = mid_channels
_base_.model.bbox_head.head_module.widen_factor = widen_factor

find_unused_parameters = True

如下表所示:参数量相近的 backbone,AttentiveNAS 提供的子网能在计算量 Flops 缩减 3 倍的情况下,与 YoloV6-s 取得近似性能。由于输出通道 channels 的缩减,Neck 部分的参数量减少,可以考虑通过 deepen_factor 扩增抵消 Neck 造成精度降低,调整 deepen_factor=1.2,此时与 Baseline(yolov6_s) 相比 mAP 提升 0.4。

尝试将输入分辨率扩增到 960,在整体总计算量小于 24.253G 的情况下,与 Baseline(yolov6_s) 相比 mAP 提升 4.1。

Table(Params; Flops)BackboneNeck + headmodelCoco mAP
Baseline (yolov6_s)12.334M; 17.22G6.535M; 7.033G18.869M; 24.253G44.0
+ attentivenas_a6 Backbone13.912M;4.41G4.081M;3.703G17.993M; 8.113G43.8 (-0.2)
+ Deepen_factor =1.213.912M, 4.41G4.463M; 4.047G18.375M, 8.487G44.4 (+0.6)
+ Image size 96013.912M, 9.915G4.463M; 9.173G18.375M; 19.088G48.1 (+3.6)

总结

本文将 MMRazor 中可搜索直接获取的 backbone 骨干网络应用与下游任务相结合,以 MMYOLO 为例,无论是 yolov5 还是 yolov6,均可通过直接替换 MMRazor 中的 backbone,实现性能的显著提升,如下表所示

Razor model vs YOLO seriesflopsparamsmap
YoloV5-S8.265G7.235M37.7
yolov5_s_spos_shufflenetv2_syncbn_8xb16-300e_coco.py7.03G7.042M38.0(+0.2)
YoloV6-S24.253G18.869M44
yolov6_l_attentive_a6_syncbn_fast_8xb16-300e_coco_channels_deeper12.py8.487G18.375M45.3(+1.3)
RTMDet-tiny8.089G4.896M41
rtmdet_tiny_ofa_lat31_syncbn_8xb32-300e_coco.py6.094G3.911M41.3(+0.3)

与此同时,在最新的 RTMDet 系列模型上,MMRazor 可搜索的 OFA 模型依然有提升效果(+0.1),对现有基线的模型参数量(-20.1%)、计算量(-24.7%)均有一定压缩。详细的 MMYOLO 配置请见: Projecs Based on MMRazor