FBNetV5——一种面向多任务的高效NAS方法

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模型的对比如下：