DepGraph：任意架構的結構化剪枝，CNN、Transformer、GNN等都適用！

1. 內容概要

本工作提出了一種非深度圖算法DepGraph，實現了架構通用的結構化剪枝，適用于CNNs, Transformers, RNNs, GNNs等網絡。該算法能夠自動地分析復雜的結構耦合，從而正確地移除參數實現網絡加速?；贒epGraph算法，我們開發了PyTorch結構化剪枝框架 Torch-Pruning。不同于依賴Masking實現的“模擬剪枝”，該框架能夠實際地移除參數和通道，降低模型推理成本。在DepGraph的幫助下，研究者和工程師無需再與復雜的網絡結構斗智斗勇，可以輕松完成復雜模型的一鍵剪枝。

論文標題：DepGraph: Towards Any Structural Pruning
論文鏈接：https://arxiv.org/abs/2301.12900
項目地址：https://github.com/VainF/Torch-Pruning

2. 背景介紹

結構化剪枝是一種重要的模型壓縮算法，它通過移除神經網絡中冗余的結構來減少參數量，從而降低模型推理的時間、空間代價。在過去幾年中，結構化剪枝技術已經被廣泛應用于各種神經網絡的加速，覆蓋了ResNet、VGG、Transformer等流行架構。然而，現有的剪枝技術依舊存在著一個棘手的問題，即算法實現和網絡結構的強綁定，這導致我們需要為不同模型分別開發專用且復雜的剪枝程序。

那么，這種強綁定從何而來？在一個網絡中，每個神經元上通常會存在多個參數連接。如下圖1（a）所示，當我們希望通過剪枝某個神經元（加粗高亮）實現加速時，與該神經元相連的多組參數需要被同步移除，這些參數就組成了結構化剪枝的最小單元，通常稱為組（Group）。然而，在不同的網絡架構中，參數的分組方式通常千差萬別。圖1（b）-（d）分別可視化了殘差結構、拼接結構、以及降維度結構所致的參數分組情況，這些結構甚至可以互相嵌套，從而產生更加復雜的分組模式。因此，參數分組也是結構化剪枝算法落地的一個難題。

圖1: 各種結構中的參數耦合，其中高亮的神經元和參數連接需要被同時剪枝

3. 本文方法

3.1 參數分組

未自動化參數分組，本文提出了一種名為DepGraph的（非深度）圖算法，來對任意網絡中的參數依賴關系進行建模。在結構化剪枝中，同一分組內的參數都是兩兩耦合的，當我們希望移除其中之一時，屬于該組的參數都需要被移除，從而保證結構的正確性。理想情況下，我們能否直接地構建一個二進制的分組矩陣G來記錄所有參數對之間耦合關系呢？如果第i層的參數和第j層參數相互耦合，我們就用來進行表示。如此，參數的分組就可以簡單建模為一個查詢問題：

然而，參數之間是否相互依賴并不僅僅由自身決定，還會受到他們之間的中間層影響。然而，中間層的結構有無窮種可能，這就導致我們難以基于規則直接判斷參數的耦合性。在分析參數依賴的過程中，我們發現一個重要的現象，即相鄰層之間的依賴關系是可以遞推的。舉個例子，相鄰的層A、B之間存在依賴，同時相鄰的層B、C之間也存在依賴，那么我們就可以遞推得到A和C之間也存在依賴關系，盡管A、C并不直接連接。這就引出了本文算法的核心，即“利用相鄰層的局部依賴關系，遞歸地推導出我們需要的分組矩陣”。而這種相鄰層間的局部依賴關系我們稱之為依賴圖（Dependency Graph），記作。依賴圖是一張稀疏且局部的關系圖，因為它僅對直接相連的層進行依賴建模。由此，分組問題可以簡化成一個路徑搜索問題，當依賴圖中節點i和節點j之間存在一條路徑時，我們可以得到，即i和j屬于同一個分組。

3.2 依賴圖建模

然而，當我們把這個簡單的想法應用到實際的網絡時，我們會發現一個新的問題。結構化剪枝中同一個層可能存在兩種剪枝方式，即輸入剪枝和輸出剪枝。對于一個卷積層而言，我們可以對參數的不同維度進行獨立的修剪，從而分別剪枝輸入通道或者輸出通道。然而，上述的依賴圖卻無法對這一現象進行建模。為此，我們提出了一種更細粒度的模型描述符，從邏輯上將每一層拆解成輸入和輸出?；谶@一描述，一個簡單的堆疊網絡就可以描述為：

其中符號表示網絡連接。還記得依賴圖是對什么關系進行建模么？答案是相鄰層的局部依賴關系！在新的模型描述方式中，“相鄰層”的定義更加廣泛，我們把同一層的輸入和輸出也視作相鄰。盡管一個神經網絡通暢包含了各種各樣的層和算子，我們依舊從上式中抽象出兩類基本依賴關系，即層間依賴（Inter-layer Dependency）和層內依賴（Intra-layer Dependency）。

層間依賴：首先我們考慮層間依賴，這種依賴關系由層和層直接的連接導致，是層類型無關的。由于一個層的輸出和下一層的輸入對應的是同一個中間特征（Feature），這就導致兩者需要被同時剪枝。例如在通道剪枝中，“某一層的的輸出通道剪枝”和“相鄰后續層的輸入通道剪枝”是等價的。

層內依賴：其次我們對層內依賴進行分析，這種依賴關系與層本身的性質有關。在神經網絡中，我們可以把各種層分為兩類：第一類層的輸入輸出可以獨立地進行剪枝，分別擁有不同的剪枝布局（pruning shcme），記作或者。例如對于全連接層的2D參數矩陣，我們可以得到和兩種不同的布局。這種情況下，輸入和輸出在依賴圖中是相互獨立、非耦合的；而另一類層輸入輸出之間存在耦合，例如逐元素運算、Batch Normalization等。他們的參數（如果有）僅有一種剪枝布局，且同時影響輸入輸出的維度。實際上，相比于復雜的參數分組模型，深度網絡中的層類型是非常有限的，我們可以預先定義不同層的剪枝布局來確定圖中的依賴關系。

綜上所述，依賴圖的構建可以基于兩條簡潔的規則實現，形式化描述為：

其中和分別表示邏輯”OR“和“AND”。我們在算法1和算法2中總結了依賴圖構建和參數分組的過程，其中參數分組是一個遞歸的連通分量（Connected Component）搜索，可以通過簡單深度或者寬度有限搜索實現。

將上述算法應用于一個具體的殘差結構塊，我們可以得到如下可視化結果。在具體剪枝時，我們以任意一個節點作為起始點，例如以作為起點，遞歸地搜索能夠訪問到的所有其他節點，并將它們歸入同一個組進行剪枝。值得注意的是，卷積網絡由于輸入輸出使用了不同的剪枝布局（），在深度圖中其輸入輸出節點間不存在依賴。而其他層例如Batch Normalization則存在依賴。

圖2: 殘差結構的依賴圖建模

3.3 利用依賴圖進行剪枝

圖3: 不同的稀疏性圖示。本方法根據依賴關系對耦合參數進行同步稀疏，從而確保剪枝掉的參數是一致“冗余”的

依賴圖的一個重要作用是參數自動分組，從而實現任意架構的模型剪枝。實際上，依賴圖的自動分組能力還可以幫助設計組級別剪枝（Group-level Pruning）。在結構化剪枝中，屬于同于組的參數會被同時移除，這一情況下我們需要保證這些被移除參數是“一致冗余”的。然而，一個常規訓練的網絡顯然不能滿足這一要求。這就需要我們引入稀疏學習方法來對參數進行稀疏化。這里同樣存在一個問題，常規的逐層獨立的稀疏技術實際上是無法實現這一目標，因為逐層算法并不考慮層間依賴關系，從而導致圖2 (b)中非一致稀疏的情況。為了解決這一問題，我們按照依賴關系將參數進行打包，如圖2 (c)所示，進行一致的稀疏訓練（虛線框內參數被推向0），從而使得耦合的參數呈現一致的重要性。在具體技術上，我們采用了一個簡單的L2正則項，通過賦予參數組的不同正則權重來進行組稀疏化。

其中k用于可剪枝參數的切片（Slicing），用于定位當前參數內第k組參數子矩陣，上述稀疏算法會得到k組不同程度稀疏的耦合參數，我們選擇整體L2 norm最小的耦合參數進行剪枝。實際上，依賴圖還可以用于設計各種更強大的組剪枝方法，但由于稀疏訓練、重要性評估等技術并非本文主要內容，這里也就不再贅述。

4 實驗

4.1 Benchmark

本文實驗主要包含兩部分，第一部分對喜聞樂見的CIFAR數據集和ImageNet數據集進行測試，我們驗證了多種模型的結構化剪枝效果，我們利用DepGraph和一致性稀疏構建了一個非常簡單的剪枝器，能夠在這兩種數據集上取得不錯的性能。

4.2 分析實驗

一致性稀疏：在分析實驗中，首先我們首先評估了一致性稀疏和逐層獨立稀疏的差異，結論符合3.3中的分析，即逐層算法無法實現依賴參數的一致稀疏。例如下圖中綠色的直方圖表示傳統的逐層稀疏策略，相比于本文提出的一致性稀疏，其整體稀疏性表現欠佳。

分組策略/稀疏度分配：我們同樣對分組策略進行了評估，我們考慮了無分組（No Grouping）、卷積分組（Conv-only）和全分組（Full Grouping）三種策略，無分組對參數進行獨立稀疏，卷積分組近考慮了卷積層而忽略其他參數化的層，最后的全分組將所有參數化的層進行一致性稀疏。實驗表明全稀疏在得到更優的結果同時，剪枝的穩定性更高，不容易出現過度剪枝的情況（性能顯著下降）。

另外剪枝的稀疏度如何分配也是一個重要問題，我們測試了算法在逐層相同稀疏度（Uniform Sparsity）和可學習稀疏度（Learned Sparsity）下的表現?？蓪W習稀疏度根據稀疏后的參數L2 Norm進行全局排序，從而決定稀疏度，這一方法能夠假設參數冗余并不是平均分布在所有層的，因此可以取得更好的性能。但于此同時，可學習的稀疏度存在過度剪枝風險，即在某一層中移除過多的參數。

依賴圖可視化：下圖中我們可視化了DenseNet-121、ResNet-18、ViT-Base的依賴圖和遞歸推導得到的分組矩陣，可以發現不同網絡的參數依賴關系是復雜且各不相同的。

非圖像模型結構化剪枝：深度模型不僅僅只有CNN和transformer，我們還對其他架構的深度模型進行了初步驗證，包括用于文本分類的LSTM，用于3D點云分類的DGCNN以及用于圖數據的GAT，我們的方法都取得了令人滿意的結果。

5 Talk is Cheap

5.1 A Minimal Example

在本節中我們展示了DepGraph的一個最簡例子。這里我們希望對一個標準ResNet-18的第一層進行通道剪枝：

通過調用DG.get_pruning_group我們可以獲取包含model.conv1的最小剪枝單位pruning_group，然后通過調用pruning_group.prune()來實現按組剪枝。通過打印這一分組，我們可以看到model.conv1上的簡單操所導致的復雜耦合：

此時，如果不依賴DepGraph，我們則需要手動進行逐層修剪，然而這通常要求開發者對網絡結構非常熟悉，同時需要手工對依賴進行分析實現分組。

5.2 High-level Pruners

基于DepGraph，我們在項目中支持了更簡單的剪枝器，用于任意架構的一鍵剪枝，目前我們已經支持了常規的權重剪枝（MagnitudePruner）、BN剪枝（BNScalePruner）、本文使用的組剪枝（GroupNormPruner）、隨機剪枝（RandomPruner）等。利用DepGraph，這些剪枝器可以快速應用到不同的模型，降低開發成本。

6 總結

本文提出了一種面向任意架構的結構化剪枝技術DepGraph，極大簡化了剪枝的流程。目前，我們的框架已經覆蓋了Torchvision模型庫中85%的模型，涵蓋分類、分割、檢測等任務?？傮w而言，本文工作只能作為“任意結架構剪枝”這一問題的初步探索性工作，無論在工程上還是在算法設計上都存在很大的改進空間。此外，當前大多數剪枝算法都是針對單層設計的，我們的工作為將來“組級別剪枝”的研究提供了一些有用的基礎資源。

審核編輯 ：李倩