DS-Net：可落地的動態網路，實際加速1.62倍，快改造起來 | CVPR 2021 Oral

晓飞的算法工程笔记發表於2024-04-09

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

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

論文: Dynamic Slimmable Network

論文地址：https://arxiv.org/abs/2103.13258
論文程式碼：https://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之間交叉熵損失。