自动网络搜索NAS在语义分割上的应用二

前言：本文将介绍如何基于 ProxylessNAS 搜索 semantic segmentation 模型，最终搜索得到的模型结构可在 CPU 上达到 36 fps 的测试结果，展示自动网络搜索（NAS）在语义分割上的应用。

随着自动网络搜索（Neural Architecture Search）技术的问世，深度学习已慢慢发展到自动化设计网络结构以及超参数配置的阶段。尤其在 AI 落地的背景下，许多模型需要部署在移动端设备。依据不同设备（GPU，CPU，芯片等），不同的模型需求（latency，模型大小，FLOPs），使用 NAS 自动搜索最佳网络结构将会是一个很有潜力的方向。上一篇介绍了 NAS 的基本框架和入门必读 DARTS [1]，以及在 semantic segmentation 领域的应用。距离现在才不过几个月，NAS 论文数量明显增长：在理论研究方面，search strategy，evaluation performance 的方法看似趋于稳定，不得不提到最近 FAIR 团队的 RegNet [2]探讨了搜索空间的设计，通过大量实验把常见的设计模型的理论一一验证，我们可以根据它的结论缩小搜索空间从而提高搜索效率；在应用方面以还是 object detection 为主，也有 segmentation，reID，GAN 等领域。

NAS 算是一项新技术，但是语义分割 semantic segmentation 却是老生常谈。自 FCN 的问世起，SegNet, UNet 这种简单粗暴的 encoder-decoder 结构在多种图像上都能达到可以接受的结果，deeplab 系列之后更是在开源数据集达到巅峰。从学术角度看 semantic segmentation 似乎已达到瓶颈，于是 researcher 们纷纷转向小样本，semi-supervised，domain adaption, cloud point 等方向另辟蹊径。但是 semantic segmentation 落地却非常困难。在实际落地场景中，使用常见的 backbone (resnet 或 yolo 系列) 就能够完成各种 object detection 任务，但是在 segmentation 上效果却不好：

1. 由于光线等原因，实际场景图像的 intensity 分布更复杂，而 segmentation 需要细分边界，对像素值的判定尤为重要。然而，相比 detection 来说 segmentation 的数据标记成本高导致训练数据较少，只依靠 data augmentation 等手段提升有限。
2. Segmentation 是 pixelwise 的任务，因为它要处理到每一个 pixel，所以模型一般都会比 object detection 的模型大许多（你看这个模型它又长又宽）。如果你的模型被要求 real-time 推理（>16 fps），那么准确度和速度必然会成为冲突，Double kill！
3. 当语义分割用在了视频流，对准确度的要求会更高。即使每两帧只相差几个 pixel，即使在 mIoU 的数值上相差无几，但是人眼看上去不够稳定，会有“抖动”的边界，Triple kill！
4. 当语义分割模型走下云端，部署在算力有限的移动端，底层芯片可能对很多操作不支持，使得在原本在可以在 GPU 上开心玩耍的模型到了 CPU 上便一朝打回解放前，Quadra kill！

Semantic segmentation 落地必须要平衡模型的准确度和速度，而设计这样的网络结构又十分困难。尝试了 BiSeNet [3]，ShuffleNetv2 [4]，MobileNetv3 [5]等一系列小模型，但是准确度和速度都没达到要求。正所谓万丈高楼平地起，成功只能靠自己，最终还是要寄希望于 NAS 自动搜索出满足条件的模型。上篇介绍的 NAS 用在语义分割还在探索阶段，在 GPU 上运行并且尝试减小 FLOPs 或 Params。但是 FLOPs 或者参数量与模型推理速度并不是正相关，只减少参数量不能满足实时推理的要求。后来的 FasterSeg [6]看似速度惊人，其实也用了 TensorRT 来加速。本文将尝试在 CPU 上完成实时的人形分割的任务，选择 ProxylessNAS 作为 baseline 来搜索模型结构。实验结果证明了 ProxylessNAS [7]还是经得起考验的，业界良心。

1.Overview of ProxylessNAS

选择 ProxylessNAS [7]的原因不仅仅是它出自名门，代码开源，在 Cifar10 和 ImageNet 数据集的准确度能从一众 NAS 模型中脱颖而出，而且它也是比较早的考虑到了模型性能的 work（如速度，模型大小，参数量）。除此之外，与 DARTS [1]系列搜索的 DAG cell 不同，ProxylessNAS [7]的主干网络采用简单的链状结构。这种链状结构 (chained-structure) 比 DAG cell 有明显的速度优势，因为它的算子之间的连接方式比较简单。

1.1 Super-net setting

我们还是用 NAS 的基本框架来解析 ProxylessNAS [7]。

Figure 1: NAS framework

• Search Space: 在搜索空间中定义的 operation candidate 的是来自 MobileNetv2 [8]的 block，分别取不同的 kernel size（3, 5, 7）和不同的 expansion rate（3，6），再加上 identity 和 zero 操作一共 8 种 ops（c.f. Figure 1）。网络的宏观结构是一个常见的链状结构来完成 classification, 每一层都有 8 个 ops candidate（c.f. Figure 2）。正如前面提到的，算子之间太复杂的连接方式会让速度变慢，常见的小模型结构都是这种链状结构。
• Search Strategy: 搜索策略采用可微分的方法，这种搜索策略近两年很常见。虽然不及 RL 和 EA 稳定，但是可以大幅度提高搜索速度。
• Evaluation Performance: One-shot 权值共享, 也是现有最常见的 super-net 的形式。对于计算资源匮乏的团队和个人来说，这种方式能够提高搜索效率减少内存占用。

1.2 Super-net training

Super-net 的参数包含两部分：operation 本身的参数和每个 operation 的权重（在 Figure2 中记为{alpha，beta，sigma …}）。将训练数据分成两部分，一部分用来训练 super-net 里面 operations 的 weight，另一部分用来更新 ops 的权重。

• Training：每个 iteration 开始的时候，在每一层都随机激活一个 operation（c.f. the binary gate in Figure 2），将所有激活的 operation 连接起来组成一条子网络记为 subnet，通过 back propagation 来更新这条 subnet 的 weight。没有激活的 ops 不放入内存，也就是说训练的时候只有一条 subnet 在内存中，这也使得整个搜索过程可以在单卡上完成。
• Searching：每个 operation 的权重 alpha 代表它的重要程度，也就是最终被选择的概率，probability = Softmax(alphas)。换言之，搜索的过程，就是不断更新权重 alpha 的过程。和 training 一样，每个 iteration 都要随机激活一条 subnet，但是这次要让 operation 的 weight 固定，通过 back propagation 计算这条 subnet 上的 alpha。Paper 里面 Eq (4)给出了计算方式，由于 binary gate 和 probability 成正比，公式里面将 loss 对 probability 的求导转化成对 binary gate 的求导，而 loss 对 binary gate 的导数在 back propagation 的时候有计算过并且保存了下来（这部分 paper 没有细说可参考源代码）。

Figure 2 illustrates the architecture of the super-net: the chained-structure searchable backbone (left) and each layer of the searchable backbone (right).

Figure 2 所表达的 ProxylessNAS 的流程，其实就是一边训练 operation 参数，一边更新 operation 的权重 alpha，最后用 Softmax 选择每一层中拥有最大 probability 的 operation 即可。读过 paper 之后确实发现有许多值得借鉴之处，但是同样也有一些疑问 (c.f. Table 1)。

Table 1 discusses the advantages and remaining issues of ProxylessNAS

2.Real-time Semantic Segmentation using ProxylessNAS on CPU

尽管对 ProxylessNAS 还有很多没有解决的问题，奈何单卡搜索训练省时省力瑕不掩瑜。借助 Intel 的 openvino 推理框架，本文尝试用 ProxylessNAS 搜索可运行在 CPU(x86)上的 real-time semantic segmentation 模型做人形分割，下面会详细介绍对算法的改进和实验结果。

2.1 Super-net setting

• Search space: 在设置搜索空间的时候，秉着大力出奇迹的心态我把常用的 operation 都塞了进来，分别是 MBv3 (3×3), MBv3 (5×5), DilConv (3×3), DilConv (5×5), SepConv (3×3), SepConv (5×5), ShuffleBlock 一共 7 种 ops。其中 MBv3 是来自 MobileNetv3 [5]的基本模块，DilConv 和 SepConv 是来自 DARTS [1]的 dilated sepatable convolutions 和 separable convolutions，ShuffleBlock 是来自 ShuffleNetv2 [4]的基本模块，前面三种 operation 都设置了两种 kernel size 可以选择。在定义宏观网络结构的时候，采用 deeplabv3+ [9]的结构 (c.f. Figure 3): head + searchable backbone + ASPP + decoder。与 UNet 类似，将 encoder 的 feature map 直接”add”到 decoder，这里没有用”concatenation”是为了避免模型过“宽”使速度变慢。其中 s2, s4, s8, s16, s32 分别指 feature map 的 resolution 下降 2，4，8，16，32 倍。与 ProxylessNAS 类似，supernet 的参数包含两部分，一部分是 operation 本身的 weight，另一部分是 operation 的权重 alpha。
• Searching Strategy: 延续 ProxylessNAS 的可微分求导方式
• Evaluation Performance: One-shot 权重共享

Figure 3 illustrates the macro-architecture of our super-net (top) and the searchable backbone (bottom)

2.2 Improvement from ProxylessNAS

• Decoupling the training and searching process: 在 ProxylessNAS 中“training”和“searching”是同时轮流完成的，也就是一边训练一边搜索。我在实验的时候把“training”和“searching”彻底分开，先用 50 个 epochs 只更新 super-net 里面 operation 的参数，在训练之后，再更新 operation 的权重 alphas。这么做的原因是避免在 operation 参数不稳定的时候，某些 alpha 过大影响后面的决策。
• Consider the latency as a hard constraint: 因为模型推理速度比较重要，而且不能用简单的叠加方式计算，所以每次随机激活 subnet 的时候都要算一下这条 subnet 的推理速度，如果不符合要求（如 latency > 30ms）则重新搜索一条 subnet，这样一定程度上避免很多推理速度过慢的 operation 被选择和学习。

2.3 Experiments

Experiment setting:

• Task: 基于 CPU（x86）的实时人像分割
• DL platform: Intel openvino

https://software.intel.com/content/www/us/en/develop/tools/openvino-toolkit.html

• Dataset : >20k 张图像，一部分来自 coco/pascal 数据集中带有”person”类别的，另一部分是私有数据
• Data augmentation: random crop, cutout, random brightness/contrast adjust, random Gaussian blur/sharpen
• Searching time: 单卡 2 GPU days (K80) 包括 training 和 searching

Experimental results:

在同一网络结构下，我们用 MobileNetv3 [5]作为 backbone 进行对比，对比结果见 Table 2。

Table 2 illustrates the experimental results

从实验数据来看，MobileNetv3 [5]的参数量和 FLOPs 都比我们搜索出来的小一倍，但是在 K80 上的推理速度很相近，准确度 mIoU 差别较大。如果综合考量准确度和速度的话，我们用 ProxylessNAS [7]搜索出来的 backbone 要明显优于 MobileNetv3 [5]的 backbone。Figure 4 的实验结果可以看出当 feature 复杂一些的时候，MobileNetv3 [5]的结果要差很多

Figure 4 compares the segmentation results of our searched network and MobileNetv3

将模型转化成 openvino 可支持模式部署在 CPU (Intel Core i7-8700)上，运行速度在 27ms 每帧左右（FPS=36），结果如 Figure 5。

Figure 5 shows the segmentation results in real application scenario

是时候展示一下搜索出来的 backbone 了，长这样~ (c.f. Figure 6)

Figure 6 illustrates the searched backbone structure

3．Future work

通过实验我们看到 ProxylessNAS 搜索策略可以从 classification 迁移到 segmentation，在速度相仿的情况下，搜索出来的网络要比原本 MobileNetv3 [5]准确度提高很多。但是只限于当前的场景，不能说人工设计出来的模型就不好或一定会被取代（虽然 MobileNetv3 也是 NAS 搜出来的）。在特定场景和有特定需求的时候，用 NAS 设计网络结构确实比人工设计加上大量调参实验要更高效，在 AI 落地方面更有发展前景。本文只是初探 ProxylessNAS，后续还会有以下几个方面的探索。

• 实验结果表明 super-net 权值共享的形式有一定合理性。但是在结构搜索的时候，将每层 probability 最大的 operation 组成 subnet 作为输出结果还是有不合理之处。因为 subnet 在搜索和训练的时候具有一定的耦合性，每层的 operation 一荣俱荣一损俱损。最终将每层最佳的 operation 选出来，组合在一起的时候未必能符合预先设定的 hard constraint，这里还是有需要改进的地方，比如可以计算相邻两层 operation 的 sub-path 的权重代替每层 operation 的权重。
• ProxylessNAS 是 MIT Hansong 团队早期的 work，现在已有后续 OFA 问世（也是跪着读完的）。在 OFA 中作者彻底将 training 和 searching 分开，结合了 knowledge distillation，先训练 teacher model，然后用 NAS 的思路在 teacher model 中搜索出最佳 student model。OFA 可以理解为自动化 network pruning 或自动 distillation。如果 OFA 实验效果好，后续还会有关于 OFA 的实战经验的分享。
• Figure 5 种实际效果展示的时候，人像和背景融合的比较自然，但是语义分割归根到底是一个分类任务，边缘的 pixel“非黑即白”，如果想要和背景自然的融合，需要计算出前景的透明度 alpha matte，这里涉及到另一项背景抠图技术，和 segmentation 配合使用效果更佳。其实 Figure 5 的下图中已经看出 segmentation 没有把头发分割出来，但是在结果中却保留了下来，也是用了背景抠图的原因。Matting 除了可以优化 segmentation 结果，还可以实现切换背景(cf. Figure 7)，PS 等功能。

下一篇我会介绍一下关于背景抠图的实战经验，敬请期待。

Figure 7 shows the demo of background matting

References

[1] Liu, Hanxiao, Karen Simonyan, and Yiming Yang. “Darts: Differentiable architecture search.” ICLR (2019).

[2] Radosavovic, Ilija, et al. “Designing Network Design Spaces.” arXiv preprint arXiv:2003.13678 (2020).

[3] Yu, Changqian, et al. “Bisenet: Bilateral segmentation network for real-time semantic segmentation.” Proceedings of the European conference on computer vision (ECCV). 2018.

[4] Zhang, Xiangyu, et al. “Shufflenet: An extremely efficient convolutional neural network for mobile devices.” Proceedings of the IEEE conference on computer vision and pattern recognition (CVPR). 2018.

[5] Howard, Andrew, et al. “Searching for mobilenetv3.” Proceedings of the IEEE International Conference on Computer Vision (ICCV). 2019.

[6] Chen, Wuyang, et al. “FasterSeg: Searching for Faster Real-time Semantic Segmentation.” ICLR (2020).

[7] Cai, Han, Ligeng Zhu, and Song Han. “Proxylessnas: Direct neural architecture search on target task and hardware.” ICLR (2019).

[8] Sandler, Mark, et al. “Mobilenetv2: Inverted residuals and linear bottlenecks.” Proceedings of the IEEE conference on computer vision and pattern recognition (CVPR). 2018.

[9] Chen, Liang-Chieh, et al. “Encoder-decoder with atrous separable convolution for semantic image segmentation.” Proceedings of the European conference on computer vision (ECCV). 2018.

点击关注，第一时间了解华为云新鲜技术~