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

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

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 的选择会构建一条单路径的子网络。

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

$W_{\mathcal{A}}=\underset{W}{\operatorname{argmin}} \mathcal{L}_{\text {train }}(\mathcal{N}(\mathcal{A}, W))$

A 代表网络搜索空间，W 代表超网权重， $\mathcal{N}(\mathcal{A}, W)$ 代表超网中编码的搜索空间。

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

$a^{*}=\underset{a \in \mathcal{A}}{\operatorname{argmax}} \mathrm{ACC}_{\mathrm{val}}\left(\mathcal{N}\left(a, W_{\mathcal{A}}(a)\right)\right)$
这个过程可以理解为将已经训练好的超网络视为一个性能评估器，对各个子网的性能进行评估，这个过程可以采用不同的方式找到最优网络，论文中主要比较了基于进化算法和随机搜索算法的两种方式。其中进化算法的伪代码为：

第三步，重训练子网：在找到最优子网络之后，会重新初始化网络结构，从头开始训练最优的子网络结构。
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=&#39;all_blocks&#39;,
space_id=f&#39;stage_{stage_idx}_block_{i}&#39;,
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 &#34;python setup.py develop&#34;
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=&#39;mmcls.ImageClassifier&#39;,
backbone=dict(
      type=&#39;SearchableShuffleNetV2&#39;, widen_factor=1.0, norm_cfg=norm_cfg),
neck=dict(type=&#39;GlobalAveragePooling&#39;),
head=dict(
      type=&#39;LinearClsHead&#39;,
      num_classes=1000,
      in_channels=1024,
      loss=dict(
         type=&#39;LabelSmoothLoss&#39;,
         num_classes=1000,
         label_smooth_val=0.1,
         mode=&#39;original&#39;,
         loss_weight=1.0),
      topk=(1, 5),
),
)

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

algorithm = dict(
type=&#39;SPOS&#39;,
architecture=dict(
      type=&#39;MMClsArchitecture&#39;,
      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[name]
         # 前向传播被选择的候选操作。
         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_ = [&#39;./spos_supernet_shufflenetv2_8xb128_in1k.py&#39;]

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

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

searcher = dict(
type=&#39;EvolutionSearcher&#39;,
candidate_pool_size=50,
candidate_top_k=10,
constraints=dict(flops=330 * 1e6),
metrics=&#39;accuracy&#39;,
score_key=&#39;accuracy_top-1&#39;,
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=&#39;top-1&#39;,       # 选择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(&#39;build search...&#39;)
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(&#39;start search...&#39;)
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_ = [
&#39;./spos_supernet_shufflenetv2_8xb128_in1k.py&#39;,
]

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[space_id][&#39;chosen&#39;]
            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 呀~

		自动登录	找回密码
密码			立即注册

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

本帖子中包含更多资源