Once for All 代码分析

Once-For-All: Train One Network And Specialize it For Efficient Deployment代码分析

• 论文链接：https://arxiv.org/abs/1908.09791

• 代码链接：https://github.com/mit-han-lab/once-for-all

• 视频链接：https://youtu.be/a_OeT8MXzWI

该论文是MIT韩松组在ICLR 2020上发表的工作，Once-For-All的核心是训练只需要训练一次超网络，就能够针对不同设备、计算量级的约束（e.g. Latency，FLOPs）选择其中不同子网络进行部署。由于只需要训练一次，并且由于predictor的引入以致搜索过程的时间和计算代价可以忽略不计，OFA能够极大减少模型部署到新设备的边际成本。该超网络具有弹性的分辨率、宽度、深度、和卷积核大小，能够根据部署需求进行设置，非常类似套娃，需要什么大小的模型，就从超网super net这个最大的娃娃中选一个合适大小的娃娃出来。超网构建这一点上，OFA和其他采用不同类型的操作ops组合成超网的NAS方法不一样，有利于训练更大的网络。

宏观算法层面的讲解可以参考顶部官方的视频链接，本文主要对代码部分进行分析。如有不足之处请批评指出。

结构

首先对主要代码的核心结构进行展示,包括超网的训练和搜索部分。

• ofa

  • imagenet_classification 超网训练部分

    • elastic_nn 弹性结构

      • modules 弹性模块定义

        • dynamic_layers.py
        • dynamic_op.py

      • networks 弹性网络定义

        • ofa_mbv3.py
        • ofa_proxyless.py
        • ofa_resnets.py

      • training 渐进式训练代码

        • progressive_shrinking.py PS train核心代码

    • networks 正常(非弹性)网络

      • mobilenet_v3.py
      • proxyless_nets.py
      • resnets.py

    • run_manager 包装网络,提供训练时需要的一些操作

  • nas 搜索部分
    • accuracy_predictor acc predictor
    • efficiency_predictor Latency 的 predictor
    • search_algorithm 搜索算法核心代码

• train_ofa_net.py fine-tune OFA超网络（progressive shrinking ）

细节

对于OFA的代码细节，主要探讨以下问题：

问题一：OFA的主要步骤是什么 ? 4个维度之间（分辨率、卷积核大小、宽度、深度）是如何安排渐进式PS训练的？

问题二：弹性网络是如何构造的？

问题三：搜索出一个好的子网之后，如何得到目标子网？子网如何创建？权重如何从超网中继承？

以下将分别讨论这三个问题。

问题一：OFA主要步骤

主要步骤

OFA主要包括三个步骤：

• 训练：正常训练一个最大的超网（动态分辨率、最大卷积核、最大宽度、最大深度）。OFA提供了官方训练的模型权重：ofa_D4_E6_K7，可直接使用。
• fine-tune：对超网ofa_D4_E6_K7进行渐进式收缩（PS）训练，尽量使得所有子网也能有很好的表现。
• 搜索：在超网中，使用进化算法和训练好的predictor在给定约束下(e.g. Latency FLOPs)进行搜索，得到满足约束的子网。

fine-tune：渐进式收缩训练（progressive shrinking ）

先定义一个采样空间，不断进行采样出子网的配置，利用该配置对超网进行设置，使其子网能够训练更新。采样方法简单采用了random.choice()进行采样。另外还设置了一个教师网络对超网进行知识蒸馏，无论是训练到哪个阶段，教师网络都是ofa_D4_E6_K7。

这4个维度的采样空间如下：

分辨率：{128 to 224 with stride 4}

卷积核大小：{3，5，7}

宽度：{3，4，6}

深度：{2，3，4}

PS训练过程

train_ofa_net.py就是用于PS训练的脚本。train_ofa_net.py是分步骤分阶段进行的，每一次执行train_ofa_net.py都会将向采样空间添加新的元素进行训练，得到不同阶段的模型权重。训练过程如下：

• 将{3，5}的卷积核大小加入采样空间进行训练，得到权重ofa_D4_E6_K357
• 将{3}的深度大小加入采样空间进行训练，得到权重ofa_D34_E6_K357
• 将{2}的深度大小加入采样空间进行训练，得到权重ofa_D234_E6_K357
• 将{4}的宽度大小加入采样空间进行训练，得到权重ofa_D234_E46_K357
• 将{3}的宽度大小加入采样空间进行训练，得到权重ofa_D234_E346_K357

ofa_D234_E346_K357即为最终PS训练完成的模型权重。

问题二：弹性网络构造

OFA超网（ofa_mbv3.py中定义的OFAMobileNetV3）主要基于MobilenetV3的Block（dynamic_layers中定义的DynamicMBConvLayer加一个短接层）进行实现。模型一共有5个units，每个units都包含depth个blocks。

在训练好一个最大卷积核kernel size=7、宽度width=6、深度 depth=4的网络之后，通过采样，得到每个block不同的设置。得到设置后，通过OFAMobileNetV3中set_active_subnet（）来对超网进行设置。设置后的网络会影响到block的forward过程，通过forward的动态化设置使得某个子网进行训练和更新而不是整个超网。重要代码部分如下：

OFAMobileNetV3中设置子网

def set_active_subnet(self, ks=None, e=None, d=None, **kwargs):
		ks = val2list(ks, len(self.blocks) - 1)
		expand_ratio = val2list(e, len(self.blocks) - 1)
		depth = val2list(d, len(self.block_group_info))

		for block, k, e in zip(self.blocks[1:], ks, expand_ratio):
			if k is not None:
				block.conv.active_kernel_size = k#修改kernel size
			if e is not None:
				block.conv.active_expand_ratio = e#修改expand ratio。MobilenetV3的宽度是以expand ratio来控制的。

		for i, d in enumerate(depth):
			if d is not None:
				self.runtime_depth[i] = min(len(self.block_group_info[i]), d)#修改网络深度

DynamicMBConvLayer 中forward动态实现

#不同layers的处理方式并不是一样的。对于DynamicMBConvLayer，只有中间的depth_conv是做了kernelsize的变化。
def forward(self, x):
        in_channel = x.size(1)

        if self.inverted_bottleneck is not None:
            self.inverted_bottleneck.conv.active_out_channel = \
                make_divisible(round(in_channel * self.active_expand_ratio), MyNetwork.CHANNEL_DIVISIBLE)#inverted_bottleneck进行expand_ratio动态设置，kernel_size不变，固定1x1

        self.depth_conv