经典网络结构搜索算法 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]