huaweiyun

前言

本文将介绍如何基于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 (3x3), MBv3 (5x5), DilConv (3x3), DilConv (5x5), SepConv (3x3), SepConv (5x5), 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:
  • 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.
 
 

相关文章: