看看谷歌如何在目標檢測任務使用預訓練權值 | CVPR 2022

論文提出能夠適配硬件加速的動態網絡DS-Net，通過提出的double-headed動態門控來實現動態路由。基於論文提出的高性能網絡設計和IEB、SGS訓練策略，僅用1/2-1/4的計算量就能達到靜態SOTA網絡性能，實際加速也有1.62倍

來源：曉飛的算法工程筆記公眾號

論文: Dynamic Slimmable Network

論文地址：//arxiv.org/abs/2103.13258
論文代碼：//github.com/changlin31/DS-Net

Introduction

模型速度在模型的移動端應用中十分重要，提高模型推理速度的方法有模型剪枝、權值量化、知識蒸餾、模型設計以及動態推理等。其中，動態推理根據輸入調整其結構，降低整體計算耗時，包含動態深度和動態維度兩個方向。如圖2所示，動態網絡自動在準確率和計算量之間trade-off，比靜態的模型設計和剪枝方法要靈活。

然而，論文發現包含動態維度的網絡的實際運行速度大都不符合預期，主要原因在於動態剪枝後的稀疏卷積與當前硬件的計算加速不匹配。大多數卷積核的動態剪枝通過zero masking(常規卷積後再通過mask取對應的輸出)或path indexing(直接通過[:,:]獲取對應的新卷積再計算)來實現，如表1所示，這些方法的計算效率都不高，導致整體推理速度沒有加快。
為了解決這一問題，論文提出了動態可精簡網絡DS-Net，在實現動態網絡的同時也有很好的硬件匹配性。
論文的主要貢獻如下：

提出新的動態網絡路由機制，通過提出的double-headed動態門控來實現網絡結構的動態路由。另外，卷積的動態剪枝通過切片的方式保持權值的內存連續性，可以很好地適配硬件加速。
提出用於DS-Net的兩階段訓練方式，包含IEB和SGS方法。IEB用於穩定可精簡網絡的訓練，SGS用於提高門控輸出的多樣性，兩者都能幫助提高DS-Net的性能。
通過ImageNet實驗對比，DS-Net的整體性能比SOTA動態網絡高約5.9%，比ResNet和MobileNet等靜態網絡性能稍微下降，但是有2-4倍計算量節省以及1.62倍實際推理加速。

Dynamic Slimmable Network

論文提出的dynamic slimmable network通過學習可精簡的超網(supernet)以及動態門控(gating)機制，達到根據不同輸入樣本動態生成網絡的目的。如圖3所示，DS-Net的超網為包含全部完整卷積的完整網絡。動態門控則是一系列預測模塊，根據輸入動態設定每個階段的卷積維度，進而生成子網，這一過程也稱為動態路由(dynamic routing)。
目前的動態網絡研究中，主網絡和動態路由通常是聯合訓練的，類似於聯合優化的網絡搜索方法。參考one-shot NAS方法，論文提出解耦的兩階段訓練方法來保證DS-Net中每個路徑的泛化性。在stage I中，禁用門控的功能並用IEB方法訓練超網，在stage II中，固定超網的權值單獨用SGS方法訓練門控。

Dynamic Supernet

這裡先介紹可在硬件高效運行的通道切片方法以及論文設計超網，然後再介紹Stage I中用到的IEB方法。

Supernet and Dynamic Channel Slicing

在如動態裁剪、動態卷積等動態網絡中，卷積核\mathcal{W}根據輸入\mathcal{X}進行動態參數化\mathcal{A}(\theta, \mathcal{X})，這樣的卷積可表示為：

動態卷積根據輸入去掉不重要的特徵通道，降低理論計算量，但其實際加速大都不符合預期。由於通道的稀疏性與硬件加速技術不匹配，在計算時不得不多次索引和拷貝需要的權值到新的連續內存空間再進行矩陣相乘。為了更好地加速，卷積核在動態權值選擇時必須保持連續且相對靜態。
基於上面的分析，論文設計了結構路由器\mathcal{A}(\theta)，能夠偏向於輸出稠密的選擇結果。對於N輸出、M輸入的卷積核W\in\mathbb{R}^{N\times M}，結構路由器輸出精簡比例\rho\in(0,1]，通過切片操作[:]選擇卷積核的前\rho\times N部分構成切片動態卷積：

[:]切片操作加*稠密矩陣乘法要比索引操作或稀疏矩陣相乘要高效得多，保證了實際運行時的速度。

SuperNet

將多個動態卷積組合起來即可搭建超網，超網通過設置不同的特徵維度組合創建多個子網。將結構路由器禁用時，超網等同於常見可精簡網絡，可用類似的方法進行預訓練。

In-place Ensemble Bootstrapping

經典的Universally Slimmable Networks通過兩個方法來有效地提升整體的性能：

sandwich rule：每次訓練的網絡組合包含最大的子網、最小的子網以及其它子網，其中最大的子網和最小的子網分別決定了可精簡網絡性能的上界和下界。
in-plcae distillation：將最大子網的向量輸出作為其它子網的訓練目標，而最大子網的訓練目標則是數據集標籤，這樣對可精簡網絡更好地收斂有很好的幫助。

雖然in-place distillation很有效，但最大子網權值的劇烈抖動會導致訓練難以收斂。根據BigNas的實驗，使用in-place distillation訓練較為複雜的網絡會極其不穩定。如果沒有殘差連接或特殊的權值初始化，在訓練初期甚至會出現梯度爆炸的情況。為了解決可精簡網絡收斂難的問題並且提升整體性能，論文提出了In-plcae Ensemble Boostrapping(IEB)方法。
首先，參考BYOL等自監督和半監督方法，使用過往的表達能力進行自監督的in-plcae distillation訓練的做法，將模型的指數滑動平均(EMA, exponential moving average)作為目標網絡生成目標向量。定義\theta和\theta^{‘}為在線網絡和目標網絡：

\alpha為動量因子，控制歷史參數的比例，t為訓練輪次。在訓練時，模型的EMA會比在線網絡更加穩定和準確，為精簡子網提供高質量的訓練目標。
接着，參考MealV2使用一組teacher網絡來生成更多樣的輸出向量供student網絡學習的做法，在進行in-place distillation時使用不同的子網構成一組teacher網絡，主要提供目標向量給最小子網學習。

整體訓練過程如圖4所示。結合sandwich rule和上述優化的in-place distillation，每論訓練有以下3種網絡：

最大的子網L使用數據集標籤作為訓練目標。
n個隨機維度的子網使用目標網絡的最大子網的向量輸出作為訓練目標。
最小的子網使用上述子網在目標網絡中對應的子網的向量輸出的組合作為訓練目標，即訓練目標為：

總結起來，超網訓練的IEB損失為：

Dynamic Slimming Gate

這裡先介紹公式2中輸出\rho因子的結構路由器\mathcal{A}(\theta, \mathcal{X})以及動態門控的double-headed設計，最後再介紹Stage II訓練使用的sandwich gate sparsification(SGS)方法。

Double-headed Design

將特徵圖轉換為精簡比例\rho有兩種方法：1）標量模式：直接通過sigmoid輸出0到1的標量作為精簡比例。2）one-hot模式：通過argmax/softmax得到one-hot向量，選擇離散的候選向量L_p中對應的精簡比例。
論文對這兩種方法進行對比後，選擇了性能更好的one-hot模式。為了將特徵圖\mathcal{X}轉換為one-hot向量，將\mathcal{A(\theta, \mathcal{X})}轉換為兩個函數的組合：

\mathcal{E}將特徵圖下採樣為向量，\mathcal{F}將向量轉化為one-hot向量用於後續的維度切片。參考DenseNet等網絡，\mathcal{E}為全局池化層，\mathcal{F}為全連接層W_1\in\mathbb{R}^{d\times C_n}+ReLU+W_2\in\mathbb{R}^{g\times d}+argmax函數(d為中間特徵維度，g為L_p的長度)：

以圖3的第n個門控為例，將大小為\rho_{n-1}C_n\times H_n\times W_n的特徵圖\mathcal{X}轉換成向量\mathcal{X}_{\mathcal{E}}\in \mathbb{R}^{\rho_{n-1}C_n}，隨後用argmax將向量進一步轉換成one-hot向量，最後通過計算one-hot向量與L_p的點積得到預測的精簡比例：

論文採用的精簡比例生成方法跟通道注意力方法十分類似，通過添加第三個全連接層W_3^{\rho_{n-1}\times d}，可直接為網絡引入注意力機制。基於上面的結構，論文提出double-headed dynamic gate，包含用於通道路由的hard channel slimming head以及用於通道注意力的soft channel attention head，其中soft channel attention head定義為：

\delta(x)=1+tanh(x)，channel attention head參與stage I的訓練。

Sandwich Gate Sparsification

在stage II訓練中，論文使用分類交叉熵損失L_{cls}和複雜度懲罰函數L_{cplx}來端到端地訓練門控，引導門控為每個輸入圖片選擇最高效的子網。為了能夠用L_{cls}來訓練不可微的slimming head，論文嘗試了經典的gumbel-softmax方法，但在實驗中發現門控很容易收斂到靜態的選項，即使加了Gumbel噪聲也優化不了。
為了解決收斂問題並且增加門控的多樣性，論文提出Sandwich Gate Sparsification(SGS)訓練方法，使用最大子網和最小子網識別輸入圖片中的hard和easy，為其生成slimming head輸出精簡因子的GT。基於訓練好的超網，將輸入大致地分為三個級別：

Easy samples \mathcal{X}_{easy}：能夠被最小子網識別的輸入。
Hard samples \mathcal{X}_{hard}：不能被最大子網識別的輸入。
Dependent samples \mathcal{X}_{dep}：不屬於上述兩種的輸入。

為了最小化計算消耗，easy samples應該都使用最小子網進行識別，即門控的GT為\mathcal{T}(\mathcal{X}_{easy})=[1,0,\cdots,0]。而對於dependent samples和hard samples則應該鼓勵其盡量使用最大的子網進行識別，即門控的GT為\mathcal{T}(\mathcal{X}_{hard})=\mathcal{T}(\mathcal{X}_{dep})=[0,0,\cdots,1]。基於這些生成的門控GT，SGS損失定義為：

\mathbb{T}_{sim}(\mathcal{X})\in{0,1}代表\mathcal{X}是否應該被最小子網預測，\mathcal{L}_{CE}(\mathcal{X},\mathcal{T})=-\sum\mathcal{T}*log(\mathcal{X})為門控輸出與生成GT之間交叉熵損失。