Paper地址:https://arxiv.org/pdf/2111.10007
前言
网络结构搜索(NAS:Neural Architecture Search),主要的Motivation有两点:
- 针对特定任务与部署平台,设计精度与速度折中最佳的模型结构;
- 实现结构设计或搜索的自动化,减轻人工成本与计算资源开销,以提升生产效率;
但影响NAS方法可泛化、普遍应用的主要原因包括:
- 当前NAS工作主要聚焦于图像分类任务,并认为适用于图像分类的最佳模型,也能成为其他任务(如检测、分割等)的最佳Backbone。然而这一假设有失偏颇,容易导致迁移应用的次优化;
- 有些NAS方法仅优化任务相关的组件,而对其他任务没有助益。例如NAS-FCOS仅对该网络涉及的FPN等组件进行结构优化,但较难应用于其他任务;
- 主流NAS方法都采用Proxyless方式(例如ProxylessNAS、FBNetV1等),需要将NAS集成至目标任务的训练Pipeline当中,因而需要花费较大的Engineering effort。当NAS方法涉及Supernet预训练时,则在不同的目标任务上执行参数共享式训练,存在一定的转换成本;
FBNetV5是沿袭FBNet系列的NAS方法,仅需一次搜索,便能够同时为不同的任务搜索相应的最佳网络结构:
- 针对多任务搜索,主要面向图像分类、目标检测与语义分割这三个CV任务,基于FBNetV3模型构建了超网络。该超网络包含多个并行分支,能够满足多种分辨率、多种网络结构(Vanilla、FPN、UNet等类型)的搜索,以适应不同CV任务对特征计算/抽象的实际需求;
- 为了将NAS从目标任务的训练Pipeline中解耦出来,设计了Multi-task proxy training,以同时满足图像分类、目标检测与语义分割等多种任务的Proxy训练;
- 设计了一种超网络预训练算法,仅需运行一次,便能够同时完成不同任务的最佳结构搜索,且搜索开销与任务数(Task number)无关;
NAS流程完成之后,需要在各自任务相关的数据集上,微调训练搜索获得的模型、以确保最佳精度。在ImageNet图像分类任务上,获得了与FBNetV3相近计算量、但精度高1.3%的模型结构;在ADE20K语义分割任务上,相比于SegFormer,所得模型的mIoU提升1.8%,计算量节省3.6x;在COCO目标检测任务上,相比于YOLOX,所得模型的mAP提升1.1%,计算量节省1.2x。
方法
下图显示了FBNetV5的搜索空间设计与搜索训练流程,具备如下优势:
- 简洁、功能丰富、可迁移的搜索空间;
- 基于Multi-task proxy,NAS流程可与目标任务训练相解耦;
- 基于所设计的超网络预训练算法,仅需运行一次搜索,便能同时生成满足不同任务需求的网络结构;
搜索空间
基于FBNetV3-A模型,扩展了模型结构以构建超网络(Supernet),其中包含多个并行分支(Parallel paths)与多个阶段(Multiple stages),每个分支的特征图分辨率(Resolution)不同,但相同阶段的分辨率相同。第一个阶段为FBNetV3模型,并划分成4个部分(4 Partitions),降采样率分别为4、8、16与32倍。对于剩余的阶段,Building block构建为Partition的最后两个网络层(Last 2 layers of the partition)。当超网络包含P个分支、S个阶段时,待选择的Block choices有B=S*P个。另外,不同阶段之间会插入融合模块(Fusion modules)以融合来自不同Path的信息,并用Binary mask确定分支Block的连接或跳过。其中Stage0~1、以及Fusion module的结构示意如下:
所构建的搜索空间(Search space)具备如下三个优势:
- 简洁、优雅:在Block-wise通过Binary mask,在B个候选Blocks之间选择最佳数据计算路径;
- 功能丰富:能够满足多种CV任务的应用需求:
- 针对图像分类任务,可应用简单的、以高分辨率分支为主的线性拓扑结构,确保高级抽象语义的表征能力;
- 针对语义分割任务,可应用UNet与PANet拓扑结构,实现Encoder-decoder架构;
- 针对目标检测任务,可应用FPN结构,以融合高级抽象语义与低级纹理信息;
- 可迁移:搜索的模型结构具备可迁移性,以适应不同模型大小、多种分辨率的应用需求;
多任务预训练搜索
为了将NAS从目标任务的训练Pipeline中解耦开来,设计了Multi-task proxy training方法,以便于将NAS快速应用于不同的任务场景。并基于ImageNet与开源模型(本文为检测与分割模型,亦可借助其他模型扩展任务应用,从而确保Scalable),构建了Pseudo-labeled多任务数据集:
- 针对分类任务,保留ImageNet的Hard label;
- 针对检测任务,通过CenterNet2为ImageNet生成目标实例的检测标签,并过滤Confidence低于0.5的检测结果;
- 针对分割任务,通过MaskFormer生成Pixel-wise的分割标签,并将最大像素概率低于0.5的分割结果设置为"Don't care category";
搜索算法
FBNetV5基于可微分搜索策略(确保充分预训练与搜索效率;可微分搜索尽管搜索空间受限,但大、小网络并排训练,有助于增强预训练效果),设计了高效率的搜索算法(与任务数无关)。如下所示,从Algorithm-1到Algorithm-4算法改进,体现了对搜索效率的增益:
实验结果
针对不同任务应用,网络结构搜索结果、以及与其他SOTA模型的对比如下:
|