打破NAS瓶颈，新方法AIO-P跨任务预测架构性能

神经网络的性能评估 (精度、召回率、PSNR 等) 需要大量的资源和时间，是神经网络结构搜索（NAS）的主要瓶颈。早期的 NAS 方法需要大量的资源来从零训练每一个搜索到的新结构。近几年来，网络性能预测器作为一种高效的性能评估方法正在引起更多关注。

然而，当前的预测器在使用范围上受限，因为它们只能建模来自特定搜索空间的网络结构，并且只能预测新结构在特定任务上的性能。例如，训练样本只包含分类网络以及它们的精度，这样训练出来的预测器只能用于评估新网络结构在图像分类任务上的性能。

​为了打破这一边界，使预测器能够预测某一网络结构在多种任务上的性能，具备跨任务跨数据泛化能力，华为海思加拿大研究院和阿尔伯塔大学联合推出了一个基于预训练和知识注入的神经网络性能预测框架。该框架可快速评估不同结构和种类的网络在分类、检测、分割等多种不同类型 CV 任务上的性能，以用于神经网络结构搜索。研究论文已被 AAAI 2023 接收。

• 论文链接：https://arxiv.org/abs/2211.17228
• 代码链接：https://github.com/Ascend-Research/AIO-P

AIO-P（All-in-One Predictors）方法旨在将神经预测器的范围扩展到分类之外的计算机视觉任务。AIO-P 利用 K-Adapter 技术将任务相关的知识注入预测器模型，同时设计了一个基于 FLOPs（浮点操作数）的标签缩放机制来适应不同的性能指标和分布。AIO-P 使用了一种独特的伪标记方案来训练 K-Adapters，仅需几分钟即可生成新的训练样本。实验结果表明，AIO-P 展示出了强大的性能预测能力，在几个计算机视觉任务上都取得了出色的 MAE 和 SRCC 结果。此外，AIO-P 可以直接迁移并预测从未见过的网络结构的性能，可以与 NAS 配合，在保证性能不降低的前提下优化现有网络的计算量。

方法介绍

AIO-P 是一种可泛化于多任务的通用网络性能预测器。AIO-P 通过预测器预训练和特定领域知识注入实现了跨任务和跨搜索空间的性能预测能力。AIO-P 利用 K-Adapter 技术将任务相关的知识注入预测器，同时依赖于通用的计算图（CG）格式表示一个网络结构，最终使其能够支持来自不同搜索空间和任务的网络，如下图 1 所示。

图 1. AIO-P 是如何表示用于不同任务的网路结构的

此外，伪标记机制的运用使 AIO-P 能够快速生成新的训练样本用以训练 K-Adapters。为了弥合不同任务上性能度量范围之间的差距，AIO-P 提出了一种基于 FLOPs 的标签缩放方法，实现了跨任务性能建模。广泛的实验结果表明，AIO-P 能够在各种不同的 CV 任务上进行准确的性能预测，如姿态估计和分割，无需训练样本或仅需少量微调。此外，AIO-P 可以正确地对从未见过的网络结构进行性能排序，与搜索算法结合后用于优化华为面部识别网络，保持其性能不变并将 FLOPs 降低超过 13.5％。该论文已被 AAAI-23 接收并且代码已经在 GitHub 上开源。

计算机视觉网络通常由执行特征提取的 “主干” 和使用提取到的特征进行预测的 “头部” 组成。“主干” 的结构通常是基于某一种已知的网络结构设计的 (ResNet, Inception, MobileNet, ViT, UNet)，而 “头部” 是针对给定任务，如分类、姿态估计、分割等而设计的。传统的 NAS 方案会根据 “主干” 的结构手动定制搜索空间，比如已知 “主干” 是 MobileNetV3，那么搜索空间可能包含 MBConv Block 数目，每个 MBConv 的参数 (kernel size, expansion)，通道数等。然而这种定制的搜索空间不具备通用性，假如有另一个 “主干” 是基于 ResNet 设计的，则无法通过现有的 NAS 框架优化它，而是需要重新设计搜索空间。

为了解决这一问题，AIO-P 选择了从计算图层面来表示不同的网络结构，实现了对任何网络结构的统一表示。具体如图 2 所示，计算图格式允许 AIO-P 将头部和主干编码在一起来表示整网结构。这也使得 AIO-P 可以预测来自不同搜索空间（如 MobileNets 和 ResNets）的网络在各种任务上的性能。

图 2. MobileNetV3 中的 Squeeze-and-Excite 模块在计算图层面的表示 

AIO-P 中提出的预测器结构从单个 GNN 回归模型开始（图 3，绿色块），它可以预测图像分类网络的性能。为了在它的基础上加入其他 CV 任务的知识，例如检测或分割，该研究将一个 K-Adapter（图 3，橙色块）附加到原始回归模型上。K-Adapter 在新任务的样本上进行训练，而原模型权重则被冻结。因此，该研究单独训练多个 K-Adapter（图 4）来加入来自多个任务的知识。

图 3. 拥有一个 K-Adapter 的 AIO-P 预测器

图 4. 拥有多个 K-Adapter 的 AIO-P 预测器

为了进一步降低训练每个 K-Adapter 的开销，该研究提出了一种巧妙的伪标签技术。这一技术使用 Latent 采样的方案来训练能共享于不同任务间的 “头部” 模型。共享头部训练之后可以与搜索空间中的任何网络主干配对，并在 10-15 分钟内进行微调以生成伪标签（图 5）。

图 5. 训练能共享于不同任务间的 “头部” 模型

经实验证明，使用共享头部获得的伪标签与通过从零开始训练一个网络一天或更长时间获得的实际性能呈正相关，有时排序相关度系数超过 0.5 (Spearman correlation)。

除此之外，不同的任务会有不同的性能指标。这些性能指标通常有自己特定的分布区间，例如，使用了某一特定主干的分类网络在 ImageNet 上分类准确率可能约为 75％，而在 MS-COCO 物体检测任务上的 mAP 可能为 30-35％。为了考虑这些不同的区间，该研究基于标准化理念提出了一种从正态分布中理解网络性能的方法。通俗的说，如果预测值为 0，则该网络性能为平均值；如果 > 0，则为较优网络；<0 对应于较差性能，无论任务、数据集或指标，如下图 6 所示。

图 6. 如何标准化网络性能

网络的 FLOPs 与模型大小，输入数据相关，并且通常与性能呈正相关趋势。该研究使用 FLOPs 转换来增强 AIO-P 从中学习的标签。

实验及结果

该研究首先在人体姿态估计和物体检测任务上上训练 AIO-P，然后用它预测多种任务上网络结构的性能，包括姿态估计（LSP 和 MPII），检测（OD），实例分割（IS），语义分割（SS）和全景分割（PS）。即使在零样本直接迁移的情况下，使用 AIO-P 对来自于 Once-for-All（OFA）搜索空间（ProxylessNAS，MobileNetV3 和 ResNet-50）的网络在这些任务上的性能进行预测，最终预测结果达到了低于 1.0％的 MAE 和超过 0.5 的排序相关度。

此外，该研究也使用 AIO-P 预测了 TensorFlow-Slim 开源模型库中的网络的性能（例如 DeepLab 语义分割模型，ResNets，Inception nets，MobileNets 和 EfficientNets），这些网络结构可能未曾在 AIO-P 的训练样本中出现。

AIO-P 通过利用 FLOPs 转换，在 3 个 DeepLab 语义分割模型库上能够实现几乎完美的 SRCC，同时在所有 4 个分类模型库上获得正的 SRCC，以及在 EfficientNet 模型上实现 SRCC=1.0。

最后，AIO-P 的核心动机是能够将其与搜索算法配对，并将其用于优化任意网络结构，可以是独立的，不属于任何搜索空间或已知模型库的结构，甚至可以是一个用于从未训练过的任务的结构。该研究使用 AIO-P 和随机变异搜索算法来优化华为手机上使用的人脸识别（FR）模型，结果显示 AIO-P 能够在降低模型计算量 FLOPs 超过 13.5％的同时保持性能（精度（Pr）和召回率（Rc））。

感兴趣的读者可以阅读论文原文，了解更多研究细节。