xiangtingsl 发表于 2022-7-19 12:51

经典网络结构搜索算法 SPOS,快速完成模型压缩

0. 引言

Single Path One Shot(以下简称 SPOS)是一种 Neural Architecture Search(网络结构搜索,以下简称 NAS)算法,相比传统的基于强化学习、进化算法等 NAS 算法,SPOS 能够显著降低搜索代价。
MMRazor 是深度学习模型压缩算法库,支持网络结构搜索、剪枝、蒸馏等主流技术方向,为 OpenMMLab 其他算法库提供即插即用、可自由组合的模型压缩算法,让用户更简单更快速地实现模型轻量化。
本文将对 SPOS 算法原理、搜索空间、MMRazor 以及 SPOS 算法在 MMRazor 中的实现做详细的解读,干货满满。

论文链接:https://arxiv.org/abs/1904.00420
官方源码:https://github.com/megvii-model/SinglePathOneShot
MMRazor 复现:https://github.com/open-mmlab/MMRazor/tree/master/configs/nas/spos
1. SPOS 算法介绍

1.1 原理介绍


SPOS 是在 ECCV2020 提出的一种 Neural Architecture Search(NAS)算法,相比传统的基于强化学习、进化算法等 NAS 算法,SPOS 能够显著降低搜索代价。
针对传统 NAS 算法中超网权重耦合度高的问题,SPOS 提出将网络权重的训练与网络结构的搜索两个过程进行解耦。首先完成超网络权重的训练,然后再从超网络中搜索出最优的子网络架构,最后对最优子网从头开始训练。
具体运行过程分为以下三步:

第一步,超网权重训练:SPOS 提出简单的由单路径候选网络构成的超网络,如下图所示,超网中每个 block 都是可选择的,并且只能选择其中一个,这样通过每层的 Choice Block 的选择会构建一条单路径的子网络。


优化过程中,通过优化每个单路径子网的权重,完成整个超网权重的优化,即以下公式。

https://www.zhihu.com/equation?tex=W_%7B%5Cmathcal%7BA%7D%7D%3D%5Cunderset%7BW%7D%7B%5Coperatorname%7Bargmin%7D%7D+%5Cmathcal%7BL%7D_%7B%5Ctext+%7Btrain+%7D%7D%28%5Cmathcal%7BN%7D%28%5Cmathcal%7BA%7D%2C+W%29%29

A 代表网络搜索空间,W 代表超网权重, https://www.zhihu.com/equation?tex=%5Cmathcal%7BN%7D%28%5Cmathcal%7BA%7D%2C+W%29 代表超网中编码的搜索空间。

第二步,网络结构搜索:从已经训练好的超网中通过进化算法找到最优的子网络。

https://www.zhihu.com/equation?tex=a%5E%7B%2A%7D%3D%5Cunderset%7Ba+%5Cin+%5Cmathcal%7BA%7D%7D%7B%5Coperatorname%7Bargmax%7D%7D+%5Cmathrm%7BACC%7D_%7B%5Cmathrm%7Bval%7D%7D%5Cleft%28%5Cmathcal%7BN%7D%5Cleft%28a%2C+W_%7B%5Cmathcal%7BA%7D%7D%28a%29%5Cright%29%5Cright%29
这个过程可以理解为将已经训练好的超网络视为一个性能评估器,对各个子网的性能进行评估,这个过程可以采用不同的方式找到最优网络,论文中主要比较了基于进化算法和随机搜索算法的两种方式。其中进化算法的伪代码为:



第三步,重训练子网:在找到最优子网络之后,会重新初始化网络结构,从头开始训练最优的子网络结构。
1.2 搜索空间介绍


SPOS 论文中提到搜索空间比较丰富,不仅包含 choiceblock 搜索,还包含了通道搜索以及混合精度量化的搜索。目前官方源码中只给出了 choiceblock 搜索,本教程也主要关注这个部分的搜索。
SPOS 的搜索空间结构如下表所示,CB 代表 choice block,一共包含了 20 个 CB。CB 内部操作的设计主要是受 ShuffleNetv2 启发,一共提供了四种操作:

[*]Choice_3: Kernel 大小为 3
[*]Choice_5: Kernel 大小为 5
[*]Choice_7: Kernel 大小为 7
[*]Choice_x: 表示 xception 架构, 其 kernel size 为 3


2. MMRazor 简介

MMRazor 是深度学习模型压缩算法库,支持网络结构搜索、剪枝、蒸馏等主流技术方向,为 OpenMMLab 其他算法库提供即插即用、可自由组合的模型压缩算法,让用户更简单更快速实现模型轻量化。MMRazor 整体设计思想和 OpenMMLab 保持一致,能够做到支持多种算法库。组织架构方面,自底向上分别由组件层、算法层和应用层构成,如下图所示。



[*]应用层支持了包括分类(MMClassification)、检测(MMDetection)和分割(MMSegmentation)为主的多种任务类型,能够复用这些库提供的功能。
[*]算法层支持了包括网络结构搜索、剪枝、知识蒸馏在内的多种模型轻量化算法。
[*]组件层提供了算法层可能需要的不同组件,比如 NAS 需要 Mutator 和 Mutable 组件;知识蒸馏需要使用到 Distiller 组件;剪枝需要使用到 Pruner 组件。

2.1 概念简介

MMRazor中引入了一些概念来实现算法,典型的概念有 Architecture、Algorithm 和 Component,其具体解释和功能如下:

(1)Architecture
架构的概念可以理解为 model,在 Architecture 中会将 model 进行封装(warpper),用于配置一些特殊的 Backbone,比如网络结构搜索中的 SuperNet,知识蒸馏中对应 Student, 剪枝中对应待剪枝模型。

(2)Algorithm
这个部分是整个 Models 模块的入口,定位上类似于 MMDetection 的 Detector。
如下图所示,算法部分涉及到的核心组件有: NAS 的组件(Mutable, Mutator), 剪枝的组件(Pruner), 蒸馏的组件(Distiller),量化的组件(Quantizer)。






这些组件之间可以灵活组织调用,比如说实现剪枝算法 autoslim 的过程中的三明治法则(即用通道较大的模型蒸馏通道较小的模型),需要用到蒸馏相关算法,因此该算法的实现可以同时用到 Pruner 和 Distiller。

(3)Component
有了这些组成 Algorithm 的不同组件之后,需要在 Algorithm 类提供的 train_step 方法中实现具体的训练步骤(类似于 Pytorch Lightning 的 LightingModule 的功能,具体执行逻辑都在 train_step 方法中)。
由于本文重点关注 NAS 算法,因此具体组件的介绍也将重点关注 Mutable 和 Mutator 组件。 在 NAS 中,需要的候选操作都是可选择或者可变化的,比如 DARTS 中提供了 7 种候选操作供搜索,SPOS 中提供了 4 种基础 Block 供选择。MMRazor 中将该选择过程抽象出来可以得到以下两个部分,分别对应两个功能:

[*]第一部分:Mutable+PlaceHolder 具有存放各个可变对象的功能,,其中 placeholder(占位符)表示当前占位部分是可变的,在开始运行之前会执行一个映射的过程,可借助 Mutator 将 PlaceHolder 通过 mapping 表映射至某个具体操作。
[*]第二部分:Mutator 具有修改 Mutable 对象属性的功能,主要负责子网采样,PlaceHolder 转化等功能。
2.2 Mutable 介绍

介绍 Mutable 之前需要引入 PlaceHolder 对象,正如其名,其作用是起到占位符的作用,表示超网中使用 PlaceHolder 构建的操作是可变对象,在执行前会将 PlaceHolder 对象通过 Mutator 进行转化,convert 成 Mutable 对象。


如上图所示,PlaceHolder 对象中保存的信息包括:

[*]group:通过 group 参数能够控制当前 PlaceHolder 具体采用的候选操作空间,该功能能够划分不同的搜索空间。
[*]space_id:对于每个 PlaceHolder 来说是某个可变操作的关键字,可以通过该关键字访问。
[*]choices:具体的一些已经注册好的候选 op。
[*]choice_args:这些 op 对应的参数配置和设置,如 kernel size、channel 等信息。

SPOS 中这样进行初始化:
Placeholder(
    group='all_blocks',
    space_id=f'stage_{stage_idx}_block_{i}',
    choice_args=dict(
      in_channels=self.in_channels,
      out_channels=out_channels,
      stride=stride,
    ))
Mutable 是专门用于实现 NAS 可搜索模块,可变对象的功能。NAS 中的搜索空间的基类是 MutableModule 提供了搜索空间中的一些必要信息,如 space_id 是相当于每个可变操作的关键字(key),可以通过 space_id 访问到具体某个可变操作;num_chosen 表示搜索空间中可以选择的操作个数。


如上图所示,MutableModule 基类下可以划分为 MutableOP 和 MutableEdge,其功能如下:

[*]MutableOP 是最常见的 One-shot NAS 中提供某个可选操作的功能。
[*]MutableEdge(后期更名为 MutableInput))则是选择不同输入进行选择的功能,最常见的是 DARTS 中的选择不同输入的功能。
通过灵活组合以上组件可以构建出不同算法所需要的搜索空间。比如 SPOS 中采用的是 MutableOP 下的 OneShotOP;DARTS 中同时采用了 DifferentiableOP 和 DifferentiableEdge 来构建搜索空间。

2.3 Mutator 介绍


Mutator 是 NAS 实现过程中最重要的组件,如下图所示,主要有以下几个功能:

[*]convert_placeholder 方法:根据用户 config 自定义的 placeholder_mapping 操作将   PlaceHolder 转为 Mutable 对象。
[*]build_search_space 方法:根据超网中的 Mutable 对象分布完成 search_space 对象的构建。
[*]其他自定义 方法:集成的子类需要实现通过改变子网的操作。比如 SPOS 中需要使用 OneShotMutator 进行随机采样;DARTS 中需要调用 DartsMutator 构建架构参数。


3. MMRazor 中超网的构建方式

在神经网络结构搜索算法中,超网的实现至关重要。算法框架中至少需要具备以下功能:(1)搜索对象是可变化的,例如 SPOS 中需要支持不同的候选操作。(2)搜索算法能够指定选择某个候选操作的功能,例如 SPOS 中在前向传播之前需要设定执行不同的子网。下面介绍 MMRazor 中如何完成这两个功能。
MMRazor 中通过引入 Mutable 和 Mutator 对象完成以上功能:
第一步,PlaceHolder 提供占位符的功能,是用户定义的可变的位置,如下图所示,其具体生效是在调用 Mutator 中的 convert 方法,转化为 Mutable 对象之后。


通过该方法可以让超网变成可搜索对象 Mutable,从而后续与 Mutator 进一步完成 NAS 任务。
第二步,在得到 Mutable 之后还需要了解 Mutable 和 Mutator 之间的交互过程,下图中展示了 SPOS 中 OneShotMutator 随机采样子网的过程:

[*]Mutator 通过调用 sample_subnet 方法,得到 Subnet_Dict 对象
[*]然后再调用 set_subnet 方法,将刚刚生成的 Subnet_Dict 对象赋予 SuperNet,决定对应的Mutable 中具体的 Choice。


4. SPOS 在 MMRazor 中的实现


4.1 环境安装

环境安装教程请参考:https://MMRazor.readthedocs.io/en/latest/get_started.html
以 cuda11.1、pytorch1.9 为例, 首先安装 cuda、torch、mmcv 包,其中 mmcv-full 表示采用了预编译包的安装方式,还需要注意对应 cuda 以及 torch 的版本。
mmcv 安装的详细方式以及 cuda、torch、mmcv 的版本对应关系可见:https://mmcv.readthedocs.io/en/latest/get_started/installation.html#
以 torch 1.9 为例进行环境的安装:
pip install torch==1.9.0+cu111 torchvision==0.10.1+cu111 torchaudio==0.9.1 -f https://download.pytorch.org/whl/torch_stable.html

pip install mmcv-full -f https://download.openmmlab.com/mmcv/dist/cu111/torch1.9.0/index.html
安装 MMRazor(推荐使用 MIM 安装):
pip install openmim
mim install MMRazor
也可以通过 pip 直接安装: pip install MMRazor

也可以通过源码安装:
git clone https://github.com/open-mmlab/MMRazor.git
cd MMRazor
pip install -v -e .# or "python setup.py develop"
4.2 Config 介绍

由于训练 SPOS 分为三个步骤,因此分别对应三个 config:

[*]SPOS 中超网络训练。对应:spos_supernet_shufflenetv2_8xb128_in1k.py
[*]通过进化算法从超网中搜索子网。对应:spos_evolution_search_shufflenetv2_8xb2048_in1k.py
[*]重训练最优子网(retrain subnet)。对应:spos_subnet_shufflenetv2_8xb128_in1k.py
以 spos_supernet_shufflenetv2_8xb128_in1k.py 为例:
model = dict(
    type='mmcls.ImageClassifier',
    backbone=dict(
      type='SearchableShuffleNetV2', widen_factor=1.0, norm_cfg=norm_cfg),
    neck=dict(type='GlobalAveragePooling'),
    head=dict(
      type='LinearClsHead',
      num_classes=1000,
      in_channels=1024,
      loss=dict(
            type='LabelSmoothLoss',
            num_classes=1000,
            label_smooth_val=0.1,
            mode='original',
            loss_weight=1.0),
      topk=(1, 5),
    ),
)

Mutator = dict(
    type='OneShotMutator',
    placeholder_mapping=dict(
      all_blocks=dict(
            type='OneShotOP',
            choices=dict(
                shuffle_3x3=dict(
                  type='ShuffleBlock', kernel_size=3, norm_cfg=norm_cfg),
                shuffle_5x5=dict(
                  type='ShuffleBlock', kernel_size=5, norm_cfg=norm_cfg),
                shuffle_7x7=dict(
                  type='ShuffleBlock', kernel_size=7, norm_cfg=norm_cfg),
                shuffle_xception=dict(
                  type='ShuffleXception', norm_cfg=norm_cfg),
            ))))

algorithm = dict(
    type='SPOS',
    architecture=dict(
      type='MMClsArchitecture',
      model=model,
    ),
    mutator=mutator,
    distiller=None,
    retraining=False,
)
可以看到 config 中主要有 model、algorithm、mutator 三个对象,其中 algorithm 中包含了 architecture 对象,architecture 对象中则包含了 model。
在初始化 algorithm 的过程中,可以在 BaseAlgorithm 中 init 函数中看到,algorithm 中会初始化 architecture, 并且根据是否传入 mutator、pruner、distiller 来决定是否初始化这三个对象:
class BaseAlgorithm(BaseModule):
    def __init__(
      self,
      architecture,
      Mutator=None,
      pruner=None,
      distiller=None,
      init_cfg=None,
    ):
      super(BaseAlgorithm, self).__init__(init_cfg)
      # 根据config build对应对象。
      self.architecture = build_architecture(architecture)
      # algorithm中可选的三个组件。
      self._init_Mutator(Mutator)
      self._init_pruner(pruner)
      self._init_distiller(distiller)
以 pruner 初始化为例:
def _init_Mutator(self, Mutator):
    if Mutator is None:
      self.Mutator = None
      return
    self.Mutator = build_Mutator(Mutator)
    self.Mutator.prepare_from_supernet(self.architecture)
    if self.retraining:
      if isinstance(self.mutable_cfg, dict):
            self.Mutator.deploy_subnet(self.architecture, self.mutable_cfg)
            self.deployed = True
      else:
            raise NotImplementedError
可以看到 Mutator 初始化函数中,除了构建 Mutator,还通过 prepare_from_supernet 方法将 architecture 中的 PlaceHolder 转化为 Mutable 对象。

4.3 超网权重训练(Pre-training)


在完成以上准备工作后,进行第一个阶段的训练:超网权重训练。如下图所示,这个过程需要不断地从超网中采样子网,并且迭代优化子网参数,最终得到优化后的超网。




训练命令如下所示:
python ./tools/mmcls/train_mmcls.py \
configs/nas/spos/spos_supernet_shufflenetv2_8xb128_in1k.py \
--work-dir $WORK_DIR SPOS 中超网训练特点是通过随机采样的方式优化网络,每次前向训练一个 batch 的过程中会随机采样一个子网络,在这里我们调用 Mutator 的 sample_subnet() 得到一组随机采样的配置,然后通过 set_subnet 方法更新 Mutable 对象中的 choice_mask 属性:
def train_step(self, data, optimizer):
    if self.retraining:
      # 如果处于retrain阶段,子网固定,进行正常训练。
      outputs = super(SPOS, self).train_step(data, optimizer)
    else:
      # 如果处于search阶段,需要Mutator采样,并通过set_subnet方法控制mutable对象
      subnet_dict = self.mutator.sample_subnet()
      self.mutator.set_subnet(subnet_dict)
      outputs = super(SPOS, self).train_step(data, optimizer)
    return outputs
SPOS 中采用的 Mutable 对象是 OneShotOP,其 forward 过程中会根据 choice_mask 访问对应被选择的 OP:

class OneShotOP(MutableOP):
    def __init__(self, **kwargs):
      super(OneShotOP, self).__init__(**kwargs)
      assert self.num_chosen == 1

    def forward(self, x):
      outputs = list()
      # choice_mask是一个one hot向量,控制选择哪些候选操作。
      for name, chosen_bool in zip(self.full_choice_names, self.choice_mask):
            if name not in self.choice_names:
                continue
            if not chosen_bool:
                continue
            module = self.choices
            # 前向传播被选择的候选操作。
            outputs.append(module(x))

      assert len(outputs) > 0
      return sum(outputs)
4.4 网络结构搜索(Evolution search)


如下图所示,Evolution Search 过程中,首先初始化候选池(Candidate Pool), 然后从预训练好的 SuperNet 中得到 Subnet 在测试集上的结果,得到对应 Top1 分数;根据得分更新候选池的 Top k 并执行 Mutation 和 CrossOver 操作,得到最优子网的网络结构。





训练命令如下所示,这里需要用到上一步超网权重的路径 $STEP1_CKPT。
python ./tools/mmcls/search_mmcls.py \
configs/nas/spos/spos_evolution_search_shufflenetv2_8xb2048_in1k.py \
$STEP1_CKPT \
--work-dir $WORK_DIR
查看对应的 config 文件,由于在搜索子网过程中采用的是进化算法,在这里引入了 Searcher 对象。 Searcher 对象在这里被归为 Core 模块中的一部分:
_base_ = ['./spos_supernet_shufflenetv2_8xb128_in1k.py']

data = dict(
    samples_per_gpu=2048,
    workers_per_gpu=16,
)

algorithm = dict(bn_training_mode=True) # 设置bn的状态为可训练

searcher = dict(
    type='EvolutionSearcher',
    candidate_pool_size=50,
    candidate_top_k=10,
    constraints=dict(flops=330 * 1e6),
    metrics='accuracy',
    score_key='accuracy_top-1',
    max_epoch=20,
    num_mutation=25,
    num_crossover=25,
    mutate_prob=0.1)
具体 Searcher 选择的是 EvolutionSearcher:
@SEARCHERS.register_module()
class EvolutionSearcher():
    def __init__(self,
               model,                                 
               dataloader,
               candidate_pool_size=10,            # 候选池大小,决定候选者数目
               candidate_top_k=5,                   # 指定top_k的候选者用于变异
               constraints=dict(flops=330 * 1e6),   # 指定限制条件
               metrics=None,                        # 指定metrics
               metric_options=None,               # 指定metri的具体选项
               score_key='top-1',          # 选择evaluate结果的具体score_key
               max_epoch=1,                         # 指定搜索的最大epoch
               num_mutation=5,                      # 指定mutation数目
               num_crossover=5,                     # 指定crossover数目
               mutate_prob=0.1,                     # 指定mutate的概率
               **search_kwargs):

    # 前向推理
    # 应用在不同的 Codebase 时,根据相应的 Codebase 重载
    @abstractmethod
    def test_fn(self, model, dataloader):            

    # 根据test_fn的结果和metric进行评估,得到衡量结果的分数
    # 应用在不同的 Codebase 时,根据相应的 Codebase 重载
    @abstractmethod
    def evaluate(self, outputs, metrics, metric_options):

    # 检测 model 是否满足限制条件
    def check_constraints(self): -> bool

    # 更新top_k的候选者
    def update_top_k(self):

    # 进化搜索主函数, 在这里完成伪代码的整体逻辑流程
    def search(self):

程序主体的执行在 search_mmcls.py 文件中,其核心运行部分为:
logger.info('build search...')
searcher = build_searcher(
    cfg.searcher,
    default_args=dict(
      algorithm=algorithm,
      dataloader=data_loader,
      test_fn=test_fn,
      work_dir=cfg.work_dir,
      logger=logger,
      resume_from=args.resume_from))
logger.info('start search...')
searcher.search()

通过 search 方法会得到最优网络的配置文件,目前已经提供了一个例子 SPOS_SHUFFLENETV2_330M_IN1k_PAPER.yaml 配置文件(与 Github 上提供的 Checkpoints 并不匹配),该文件中保存的是每个 stage 中每个 block 具体选择的候选操作名称, 如下所示:

stage_0_block_0: # 第0个stage的第0个block
chosen:
- shuffle_7x7# 选择的候选操作为Shuffle_7x7
stage_0_block_1: # 第0个stage的第1个block
chosen:
- shuffle_5x5# 选择的候选操作为Shuffle_5x5
stage_0_block_2:
chosen:
- shuffle_3x3
stage_0_block_3:
chosen:
- shuffle_5x5

4.5 重训练子网(Retrain)


如下图所示,在上一步通过进化算法得到最优子网结构之后,将其对应的子网络从头进行训练,得到最终的可用的网络模型,整个流程与训练普通分类网络是一致的。





训练命令如下所示:这里需要将 algorithm.mutable_cfg 参数传入,该参数就是上一步得到的 yaml 文件的位置。
python ./tools/mmcls/train_mmcls.py \
configs/nas/spos/spos_subnet_shufflenetv2_8xb128_in1k.py \
--work-dir $WORK_DIR \
--cfg-options algorithm.mutable_cfg=$STEP2_SUBNET_YAML
这个部分跟 MMClassfication 训练一个普通网络的过程完全一致,使用的 config 如下,大部分都是复用了第一步:
_base_ = [
    './spos_supernet_shufflenetv2_8xb128_in1k.py',
]

algorithm = dict(retraining=True)

runner = dict(max_iters=300000)
find_unused_parameters = False
但是需要注意,除了需要设置 mutable_cfg 外,还需要将 retraining 设置为 True。
由于在网络结构搜索和剪枝过程中会分为多个阶段,train stage 和 retrain stage,在 retrain 阶段通常会选择给定的某个子网络,以 NAS 任务为例:
def _init_mutator(self, mutator):
    if mutator is None:
      self.mutator = None
      return
    self.mutator = build_mutator(mutator)
    self.mutator.prepare_from_supernet(self.architecture)
    if self.retraining:
      if isinstance(self.mutable_cfg, dict):
            self.mutator.deploy_subnet(self.architecture, self.mutable_cfg)
            self.deployed = True
      else:
            raise NotImplementedError

retraining=True 会导致调用 mutator.deploy_subnet,该方法根据传入的 mutable_cfg 固定子网络的结构:
def deploy_subnet(self, supernet, subnet_dict):
   def traverse(module):
      for name, child in module.named_children():
            if isinstance(child, MutableModule):
                space_id = child.space_id
                chosen = subnet_dict['chosen']
                child.export(chosen)
            traverse(child)
    traverse(supernet)
该方法通过调用每个 Mutable 对象的 export 方法, 能够将 choice 属性中没有被采样到的候选 OP 删除,通过这种方式得到子网络:
def export(self, chosen):
    for name in self.choice_names:
      if name not in chosen:
            self.choices.pop(name) 5. 总结

本文主要解读了经典的网络结构搜索算法 SPOS, 介绍了 MMRazor 的模型组件,以及如何在 MMRazor 中运行该算法的流程,将核心的代码分享给大家。在实际使用中,SPOS 能够和各代码库搭配使用,例如通过与 MMDetection 库的配合,能够简便地实现 DetNAS 算法。

MMRazor 不仅包含 NAS 相关算法,也有蒸馏和剪枝等功能。欢迎大家体验,如果对你有帮助的话,欢迎点个star 呀~
页: [1]
查看完整版本: 经典网络结构搜索算法 SPOS,快速完成模型压缩