DepGraph：任意架構(gòu)的結(jié)構(gòu)化剪枝，CNN、Transformer、GNN等都適用！

1. 內(nèi)容概要

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

論文標(biāo)題：DepGraph: Towards Any Structural Pruning
論文鏈接：https://arxiv.org/abs/2301.12900
項(xiàng)目地址：https://github.com/VainF/Torch-Pruning

2. 背景介紹

結(jié)構(gòu)化剪枝是一種重要的模型壓縮算法，它通過(guò)移除神經(jīng)網(wǎng)絡(luò)中冗余的結(jié)構(gòu)來(lái)減少參數(shù)量，從而降低模型推理的時(shí)間、空間代價(jià)。在過(guò)去幾年中，結(jié)構(gòu)化剪枝技術(shù)已經(jīng)被廣泛應(yīng)用于各種神經(jīng)網(wǎng)絡(luò)的加速，覆蓋了ResNet、VGG、Transformer等流行架構(gòu)。然而，現(xiàn)有的剪枝技術(shù)依舊存在著一個(gè)棘手的問(wèn)題，即算法實(shí)現(xiàn)和網(wǎng)絡(luò)結(jié)構(gòu)的強(qiáng)綁定，這導(dǎo)致我們需要為不同模型分別開(kāi)發(fā)專用且復(fù)雜的剪枝程序。

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

圖1: 各種結(jié)構(gòu)中的參數(shù)耦合，其中高亮的神經(jīng)元和參數(shù)連接需要被同時(shí)剪枝

3. 本文方法

3.1 參數(shù)分組

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

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

3.2 依賴圖建模

然而，當(dāng)我們把這個(gè)簡(jiǎn)單的想法應(yīng)用到實(shí)際的網(wǎng)絡(luò)時(shí)，我們會(huì)發(fā)現(xiàn)一個(gè)新的問(wèn)題。結(jié)構(gòu)化剪枝中同一個(gè)層可能存在兩種剪枝方式，即輸入剪枝和輸出剪枝。對(duì)于一個(gè)卷積層而言，我們可以對(duì)參數(shù)的不同維度進(jìn)行獨(dú)立的修剪，從而分別剪枝輸入通道或者輸出通道。然而，上述的依賴圖卻無(wú)法對(duì)這一現(xiàn)象進(jìn)行建模。為此，我們提出了一種更細(xì)粒度的模型描述符，從邏輯上將每一層拆解成輸入和輸出?；谶@一描述，一個(gè)簡(jiǎn)單的堆疊網(wǎng)絡(luò)就可以描述為：

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

層間依賴：首先我們考慮層間依賴，這種依賴關(guān)系由層和層直接的連接導(dǎo)致，是層類型無(wú)關(guān)的。由于一個(gè)層的輸出和下一層的輸入對(duì)應(yīng)的是同一個(gè)中間特征（Feature），這就導(dǎo)致兩者需要被同時(shí)剪枝。例如在通道剪枝中，“某一層的的輸出通道剪枝”和“相鄰后續(xù)層的輸入通道剪枝”是等價(jià)的。

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

綜上所述，依賴圖的構(gòu)建可以基于兩條簡(jiǎn)潔的規(guī)則實(shí)現(xiàn)，形式化描述為：

其中和分別表示邏輯”O(jiān)R“和“AND”。我們?cè)谒惴?和算法2中總結(jié)了依賴圖構(gòu)建和參數(shù)分組的過(guò)程，其中參數(shù)分組是一個(gè)遞歸的連通分量（Connected Component）搜索，可以通過(guò)簡(jiǎn)單深度或者寬度有限搜索實(shí)現(xiàn)。

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

圖2: 殘差結(jié)構(gòu)的依賴圖建模

3.3 利用依賴圖進(jìn)行剪枝

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

依賴圖的一個(gè)重要作用是參數(shù)自動(dòng)分組，從而實(shí)現(xiàn)任意架構(gòu)的模型剪枝。實(shí)際上，依賴圖的自動(dòng)分組能力還可以幫助設(shè)計(jì)組級(jí)別剪枝（Group-level Pruning）。在結(jié)構(gòu)化剪枝中，屬于同于組的參數(shù)會(huì)被同時(shí)移除，這一情況下我們需要保證這些被移除參數(shù)是“一致冗余”的。然而，一個(gè)常規(guī)訓(xùn)練的網(wǎng)絡(luò)顯然不能滿足這一要求。這就需要我們引入稀疏學(xué)習(xí)方法來(lái)對(duì)參數(shù)進(jìn)行稀疏化。這里同樣存在一個(gè)問(wèn)題，常規(guī)的逐層獨(dú)立的稀疏技術(shù)實(shí)際上是無(wú)法實(shí)現(xiàn)這一目標(biāo)，因?yàn)橹饘铀惴ú⒉豢紤]層間依賴關(guān)系，從而導(dǎo)致圖2 (b)中非一致稀疏的情況。為了解決這一問(wèn)題，我們按照依賴關(guān)系將參數(shù)進(jìn)行打包，如圖2 (c)所示，進(jìn)行一致的稀疏訓(xùn)練（虛線框內(nèi)參數(shù)被推向0），從而使得耦合的參數(shù)呈現(xiàn)一致的重要性。在具體技術(shù)上，我們采用了一個(gè)簡(jiǎn)單的L2正則項(xiàng)，通過(guò)賦予參數(shù)組的不同正則權(quán)重來(lái)進(jìn)行組稀疏化。

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

4 實(shí)驗(yàn)

4.1 Benchmark

本文實(shí)驗(yàn)主要包含兩部分，第一部分對(duì)喜聞樂(lè)見(jiàn)的CIFAR數(shù)據(jù)集和ImageNet數(shù)據(jù)集進(jìn)行測(cè)試，我們驗(yàn)證了多種模型的結(jié)構(gòu)化剪枝效果，我們利用DepGraph和一致性稀疏構(gòu)建了一個(gè)非常簡(jiǎn)單的剪枝器，能夠在這兩種數(shù)據(jù)集上取得不錯(cuò)的性能。

4.2 分析實(shí)驗(yàn)

一致性稀疏：在分析實(shí)驗(yàn)中，首先我們首先評(píng)估了一致性稀疏和逐層獨(dú)立稀疏的差異，結(jié)論符合3.3中的分析，即逐層算法無(wú)法實(shí)現(xiàn)依賴參數(shù)的一致稀疏。例如下圖中綠色的直方圖表示傳統(tǒng)的逐層稀疏策略，相比于本文提出的一致性稀疏，其整體稀疏性表現(xiàn)欠佳。

分組策略/稀疏度分配：我們同樣對(duì)分組策略進(jìn)行了評(píng)估，我們考慮了無(wú)分組（No Grouping）、卷積分組（Conv-only）和全分組（Full Grouping）三種策略，無(wú)分組對(duì)參數(shù)進(jìn)行獨(dú)立稀疏，卷積分組近考慮了卷積層而忽略其他參數(shù)化的層，最后的全分組將所有參數(shù)化的層進(jìn)行一致性稀疏。實(shí)驗(yàn)表明全稀疏在得到更優(yōu)的結(jié)果同時(shí)，剪枝的穩(wěn)定性更高，不容易出現(xiàn)過(guò)度剪枝的情況（性能顯著下降）。

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

依賴圖可視化：下圖中我們可視化了DenseNet-121、ResNet-18、ViT-Base的依賴圖和遞歸推導(dǎo)得到的分組矩陣，可以發(fā)現(xiàn)不同網(wǎng)絡(luò)的參數(shù)依賴關(guān)系是復(fù)雜且各不相同的。

非圖像模型結(jié)構(gòu)化剪枝：深度模型不僅僅只有CNN和transformer，我們還對(duì)其他架構(gòu)的深度模型進(jìn)行了初步驗(yàn)證，包括用于文本分類的LSTM，用于3D點(diǎn)云分類的DGCNN以及用于圖數(shù)據(jù)的GAT，我們的方法都取得了令人滿意的結(jié)果。

5 Talk is Cheap

5.1 A Minimal Example

在本節(jié)中我們展示了DepGraph的一個(gè)最簡(jiǎn)例子。這里我們希望對(duì)一個(gè)標(biāo)準(zhǔn)ResNet-18的第一層進(jìn)行通道剪枝：

通過(guò)調(diào)用DG.get_pruning_group我們可以獲取包含model.conv1的最小剪枝單位pruning_group，然后通過(guò)調(diào)用pruning_group.prune()來(lái)實(shí)現(xiàn)按組剪枝。通過(guò)打印這一分組，我們可以看到model.conv1上的簡(jiǎn)單操所導(dǎo)致的復(fù)雜耦合：

此時(shí)，如果不依賴DepGraph，我們則需要手動(dòng)進(jìn)行逐層修剪，然而這通常要求開(kāi)發(fā)者對(duì)網(wǎng)絡(luò)結(jié)構(gòu)非常熟悉，同時(shí)需要手工對(duì)依賴進(jìn)行分析實(shí)現(xiàn)分組。

5.2 High-level Pruners

基于DepGraph，我們?cè)陧?xiàng)目中支持了更簡(jiǎn)單的剪枝器，用于任意架構(gòu)的一鍵剪枝，目前我們已經(jīng)支持了常規(guī)的權(quán)重剪枝（MagnitudePruner）、BN剪枝（BNScalePruner）、本文使用的組剪枝（GroupNormPruner）、隨機(jī)剪枝（RandomPruner）等。利用DepGraph，這些剪枝器可以快速應(yīng)用到不同的模型，降低開(kāi)發(fā)成本。

6 總結(jié)

本文提出了一種面向任意架構(gòu)的結(jié)構(gòu)化剪枝技術(shù)DepGraph，極大簡(jiǎn)化了剪枝的流程。目前，我們的框架已經(jīng)覆蓋了Torchvision模型庫(kù)中85%的模型，涵蓋分類、分割、檢測(cè)等任務(wù)。總體而言，本文工作只能作為“任意結(jié)架構(gòu)剪枝”這一問(wèn)題的初步探索性工作，無(wú)論在工程上還是在算法設(shè)計(jì)上都存在很大的改進(jìn)空間。此外，當(dāng)前大多數(shù)剪枝算法都是針對(duì)單層設(shè)計(jì)的，我們的工作為將來(lái)“組級(jí)別剪枝”的研究提供了一些有用的基礎(chǔ)資源。

審核編輯 ：李倩