AAAI 2023｜打破NAS瓶頸，新方法AIO-P跨任務(wù)預測架構性能

機器之心專(zhuān)欄

機器之心編輯部

華為海思加拿大研究院和阿爾伯塔大學(xué)吉印通推出了一個(gè)基于預訓練和知識注入的神經(jīng)網(wǎng)絡(luò )性能預測框架。

神經(jīng)網(wǎng)絡(luò )的性能評估 (精度、召回率、PSNR 等) 需要大量的資源和時(shí)間，是神經(jīng)網(wǎng)絡(luò )結構搜索（NAS）的主要瓶頸。早期的 NAS 方法需要大量的資源來(lái)從零訓練每一個(gè)搜索到的新結構。近幾年來(lái)，網(wǎng)絡(luò )性能預測器作為一種高效的性能評估方法正在引起更多關(guān)注。

然而，當前的預測器在使用范圍上受限，因為它們只能建模來(lái)自特定搜索空間的網(wǎng)絡(luò )結構，并且只能預測新結構在特定任務(wù)上的性能。例如，訓練樣本只包含分類(lèi)網(wǎng)絡(luò )以及它們的精度，這樣訓練出來(lái)的預測器只能用于評估新網(wǎng)絡(luò )結構在圖像分類(lèi)任務(wù)上的性能。

為了打破這一邊界，使預測器能夠預測某一網(wǎng)絡(luò )結構在多種任務(wù)上的性能，具備跨任務(wù)跨數據泛化能力，華為海思加拿大研究院和阿爾伯塔大學(xué)吉印通推出了一個(gè)基于預訓練和知識注入的神經(jīng)網(wǎng)絡(luò )性能預測框架。該框架可快速評估不同結構和種類(lèi)的網(wǎng)絡(luò )在分類(lèi)、檢測、分割等多種不同類(lèi)型 CV 任務(wù)上的性能，以用于神經(jīng)網(wǎng)絡(luò )結構搜索。研究論文已被 AAAI 2023 接收。

論文鏈接：

代碼鏈接：

021yin.com/Ascend-Research/AIO-P

AIO-P（All-in-One Predictors）方法旨在將神經(jīng)預測器的范圍擴展到分類(lèi)之外的計算機視覺(jué)任務(wù)。AIO-P 利用 K-Adapter 技術(shù)將任務(wù)相關(guān)的知識注入預測器模型，同時(shí)設計了一個(gè)基于 FLOPs（浮點(diǎn)操作數）的標簽縮放機制來(lái)適應不同的性能指標和分布。AIO-P 使用了一種獨特的偽標記方案來(lái)訓練 K-Adapters，僅需幾分鐘即可生成新的訓練樣本。實(shí)驗結果表明，AIO-P 展示出了強大的性能預測能力，在幾個(gè)計算機視覺(jué)任務(wù)上都取得了出色的 MAE 和 SRCC 結果。此外，AIO-P 可以直接遷移并預測從未見(jiàn)過(guò)的網(wǎng)絡(luò )結構的性能，可以與 NAS 配合，在保證性能不降低的前提下優(yōu)化現有網(wǎng)絡(luò )的計算量。

方法介紹

AIO-P 是一種可泛化于多任務(wù)的通用網(wǎng)絡(luò )性能預測器。AIO-P 通過(guò)預測器預訓練和特定領(lǐng)域知識注入實(shí)現了跨任務(wù)和跨搜索空間的性能預測能力。AIO-P 利用 K-Adapter 技術(shù)將任務(wù)相關(guān)的知識注入預測器，同時(shí)依賴(lài)于通用的計算圖（CG）格式表示一個(gè)網(wǎng)絡(luò )結構，最終使其能夠支持來(lái)自不同搜索空間和任務(wù)的網(wǎng)絡(luò )，如下圖 1 所示。

圖 1. AIO-P 是如何表示用于不同任務(wù)的網(wǎng)路結構的

此外，偽標記機制的運用使 AIO-P 能夠快速生成新的訓練樣本用以訓練 K-Adapters。為了彌合不同任務(wù)上性能度量范圍之間的差距，AIO-P 提出了一種基于 FLOPs 的標簽縮放方法，實(shí)現了跨任務(wù)性能建模。廣泛的實(shí)驗結果表明，AIO-P 能夠在各種不同的 CV 任務(wù)吉印通行準確的性能預測，如姿態(tài)估計和分割，無(wú)需訓練樣本或僅需少量微調。此外，AIO-P 可以正確地對從未見(jiàn)過(guò)的網(wǎng)絡(luò )結構進(jìn)行性能排序，與搜索算法結合后用于優(yōu)化華為面部識別網(wǎng)絡(luò )，保持其性能不變并將 FLOPs 降低超過(guò) 13.5％。該論文已被 AAAI-23 接收并且代碼已經(jīng)在 GitHub 上開(kāi)源。

計算機視覺(jué)網(wǎng)絡(luò )通常由執行特征提取的 “主干” 和使用提取到的特征進(jìn)行預測的 “頭部” 組成?！爸鞲伞?的結構通常是基于某一種已知的網(wǎng)絡(luò )結構設計的 (ResNet, Inception, MobileNet, ViT, UNet)，而 “頭部” 是針對給定任務(wù)，如分類(lèi)、姿態(tài)估計、分割等而設計的。傳統的 NAS 方案會(huì )根據 “主干” 的結構手動(dòng)定制搜索空間，比如已知 “主干” 是 MobileNetV3，那么搜索空間可能包含 MBConv Block 數目，每個(gè) MBConv 的參數 (kernel size, expansion)，通道數等。然而這種定制的搜索空間不具備通用性，假如有另一個(gè) “主干” 是基于 ResNet 設計的，則無(wú)法通過(guò)現有的 NAS 框架優(yōu)化它，而是需要重新設計搜索空間。

為了解決這一問(wèn)題，AIO-P 選擇了從計算圖層面來(lái)表示不同的網(wǎng)絡(luò )結構，實(shí)現了對任何網(wǎng)絡(luò )結構的統一表示。具體如圖 2 所示，計算圖格式允許 AIO-P 將頭部和主干編碼在一起來(lái)表示整網(wǎng)結構。這也使得 AIO-P 可以預測來(lái)自不同搜索空間（如 MobileNets 和 ResNets）的網(wǎng)絡(luò )在各種任務(wù)上的性能。

圖 2. MobileNetV3 中的 Squeeze-and-Excite 模塊在計算圖層面的表示

AIO-P 中提出的預測器結構從單個(gè) GNN 回歸模型開(kāi)始（圖 3，綠色塊），它可以預測圖像分類(lèi)網(wǎng)絡(luò )的性能。為了在它的基礎上加入其他 CV 任務(wù)的知識，例如檢測或分割，該研究將一個(gè) K-Adapter（圖 3，橙色塊）附加到原始回歸模型上。K-Adapter 在新任務(wù)的樣本吉印通行訓練，而原模型權重則被凍結。因此，該研究單獨訓練多個(gè) K-Adapter（圖 4）來(lái)加入來(lái)自多個(gè)任務(wù)的知識。

圖 3. 擁有一個(gè) K-Adapter 的 AIO-P 預測器

圖 4. 擁有多個(gè) K-Adapter 的 AIO-P 預測器

為了進(jìn)一步降低訓練每個(gè) K-Adapter 的開(kāi)銷(xiāo)，該研究提出了一種巧妙的偽標簽技術(shù)。這一技術(shù)使用 Latent 采樣的方案來(lái)訓練能共享于不同任務(wù)間的 “頭部” 模型。共享頭部訓練之后可以與搜索空間中的任何網(wǎng)絡(luò )主干配對，并在 10-15 分鐘內進(jìn)行微調以生成偽標簽（圖 5）。

圖 5. 訓練能共享于不同任務(wù)間的 “頭部” 模型

經(jīng)實(shí)驗證明，使用共享頭部獲得的偽標簽與通過(guò)從零開(kāi)始訓練一個(gè)網(wǎng)絡(luò )一天或更長(cháng)時(shí)間獲得的實(shí)際性能呈正相關(guān)，有時(shí)排序相關(guān)度系數超過(guò) 0.5 (Spearman correlation)。

除此之外，不同的任務(wù)會(huì )有不同的性能指標。這些性能指標通常有自己特定的分布區間，例如，使用了某一特定主干的分類(lèi)網(wǎng)絡(luò )在 ImageNet 上分類(lèi)準確率可能約為 75％，而在 MS-COCO 物體檢測任務(wù)上的 mAP 可能為 30-35％。為了考慮這些不同的區間，該研究基于標準化理念提出了一種從正態(tài)分布中理解網(wǎng)絡(luò )性能的方法。通俗的說(shuō)，如果預測值為 0，則該網(wǎng)絡(luò )性能為平均值；如果 0，則為較優(yōu)網(wǎng)絡(luò )；

圖 6. 如何標準化網(wǎng)絡(luò )性能

網(wǎng)絡(luò )的 FLOPs 與模型大小，輸入數據相關(guān)，并且通常與性能呈正相關(guān)趨勢。該研究使用 FLOPs 轉換來(lái)增強 AIO-P 從中學(xué)習的標簽。

實(shí)驗及結果

該研究首先在人體姿態(tài)估計和物體檢測任務(wù)上上訓練 AIO-P，然后用它預測多種任務(wù)上網(wǎng)絡(luò )結構的性能，包括姿態(tài)估計（LSP 和 MPII），檢測（OD），實(shí)例分割（IS），語(yǔ)義分割（SS）和全景分割（PS）。即使在零樣本直接遷移的情況下，使用 AIO-P 對來(lái)自于 Once-for-All（OFA）搜索空間（ProxylessNAS，MobileNetV3 和 ResNet-50）的網(wǎng)絡(luò )在這些任務(wù)上的性能進(jìn)行預測，最終預測結果達到了低于 1.0％的 MAE 和超過(guò) 0.5 的排序相關(guān)度。

此外，該研究也使用 AIO-P 預測了 TensorFlow-Slim 開(kāi)源模型庫中的網(wǎng)絡(luò )的性能（例如 DeepLab 語(yǔ)義分割模型，ResNets，Inception nets，MobileNets 和 EfficientNets），這些網(wǎng)絡(luò )結構可能未曾在 AIO-P 的訓練樣本中出現。

AIO-P 通過(guò)利用 FLOPs 轉換，在 3 個(gè) DeepLab 語(yǔ)義分割模型庫上能夠實(shí)現幾乎完美的 SRCC，同時(shí)在所有 4 個(gè)分類(lèi)模型庫上獲得正的 SRCC，以及在 EfficientNet 模型上實(shí)現 SRCC=1.0。

最后，AIO-P 的核心動(dòng)機是能夠將其與搜索算法配對，并將其用于優(yōu)化任意網(wǎng)絡(luò )結構，可以是獨立的，不屬于任何搜索空間或已知模型庫的結構，甚至可以是一個(gè)用于從未訓練過(guò)的任務(wù)的結構。該研究使用 AIO-P 和隨機變異搜索算法來(lái)優(yōu)化華為手機上使用的人臉識別（FR）模型，結果顯示 AIO-P 能夠在降低模型計算量 FLOPs 超過(guò) 13.5％的同時(shí)保持性能（精度（Pr）和召回率（Rc））。

感興趣的讀者可以閱讀論文原文，了解更多研究細節。