基于模型壓縮與重構U-net的胰腺分割

馬 豪，劉 彥，張俊然

(四川大學 電氣工程學院，四川 成都 610000)

0 引 言

與其它腹部器官相比，胰腺擁有如下特性：其個體之間的形狀、大小和位置差異很大，且與周圍組織環境之間對比度低。傳統的胰腺分割算法大概可分為如下兩種模式：從分割模型架構上來講往往采用雙層分割網絡來保障胰腺的分割精度；從特征圖角度來講可對胰腺提取不同維度的特征圖信息以充分利用標簽特征。其算法目的均是為了追求提升胰腺分割精度。但是這些算法缺點在于除了受胰腺自身特性影響，也受限于對分割模型卷積層輸出的特征圖信息利用不充分，致使最終的分割精度往往不能達到醫生滿意的程度，因此對于胰腺的分割仍是一個很大的挑戰。最近由于深度學習中卷積神經網絡對于硬件的苛刻要求和醫療設備算力薄弱的原因，一批專家學者開始應用模型壓縮技術來降低卷積神經網絡的計算需求與存儲空間占用。融合了上述兩種手段的實時醫學圖像分割課題，已然成為熱門研究領域[1-6]。綜上所述，可見醫學上實時胰腺圖像分割的研究目的為：在保障分割網絡性能的前提下，盡可能減少分割模型對硬件資源的需求。針對上述問題，本文利用通道剪枝技術對分割模型進行充分壓縮，同時利用重構U-net對胰腺進行圖像分割處理。相較前人，在分割精度與分割模型存儲空間占用上同時做出新的突破，更進一步推進了對于胰腺的醫學上實時圖像分割實現。

1 相關工作

1.1 胰腺分割

對于胰腺的分割算法研究經歷了許多階段，最早期的胰腺分割算法是基于模型形狀統計設計，或者運用多圖譜技術[7,8]。隨后深度學習開始應用到胰腺分割領域上來。從對胰腺分割算法的結構設計方面講，一開始的胰腺分割算法一般做法都是應用單個分割網絡，但是由于胰腺圖像所占CT圖像比例非常小(約為0.5%)，單個分割網絡對于胰腺的特征信息提取不夠充分，導致分割精度很低，所以后續的胰腺分割算法往往采用雙層的粗細分割網絡結構來處理胰腺數據。遞歸顯著變換神經網絡提出在粗細分割網絡之間加入一個循環分割模塊，用以將粗分割網絡的權重轉移到精細分割網絡中去，使兩個分割網絡進行共同迭代優化直至訓練結束，最終取得了Dice為84%的成績[9,10]。

對于胰腺的分割網絡從維度上分類，一般分為3種：3D/2D/2.5D分割網絡。胰腺分割的一種3D分割網絡為首先給3DFCN模型加入了Resnet連接，并提取不同維度特征圖進行特征融合，再將其又分別反饋給對應位置初始的特征圖，最后使用深度監督模塊處理后得到最終分割結果，其在NIH胰腺公共數據集上取得的Dice參數為83.39%。但是問題在于3D的分割網絡十分消耗硬件資源，所以現階段2.5D分割網絡在胰腺數據集上應用更為廣泛。針對胰腺數據的長短期記憶網絡為每次經過3張連續的特征圖融合得到一張新的特征圖，最終構成分割網絡，相比于其它2.5D網絡該分割網絡所需的參數量更少[11-13]。

前人工作在只追求胰腺分割精度的前提下，目前取得的最好分割精度是由GAN網絡實現的。GAN網絡在胰腺分割上的應用方法是：用GAN網絡融合膨脹卷積以獲得更大的感受野，同時加入局部長短期記憶網絡來保持相鄰切片之間的上下文相關性,接著加入對抗模塊用全局分布約束來保持全局空間平滑一致性。該文獻在NIH數據集上最終取得的Dice分割精度為88.72%[14]。

1.2 模型壓縮

卷積神經網絡的數據壓縮方法大概可以分為4類：量化、知識蒸餾、剪枝以及緊湊網絡設計。其中最主流的兩種方法是剪枝和量化。剪枝即基于卷積神經網絡中參數的重要程度不同，可以只保留重要參數而剪去不重要的參數或通道，這樣可以將網絡進行高度稀疏化并獲得相近的卷積神經網絡效能。根據剪枝的細粒度不同，將剪枝方法分為5類：細粒度剪枝、矢量級剪枝、核級剪枝、組級剪枝和過濾級剪枝[15]。量化在卷積神經網絡的壓縮和時間加速方面有著巨大作用。對于神經網絡參數的壓縮又分為對權重壓縮和對激活壓縮。三元權值網絡提出將權重參數約束在[-1，0，1]中取值，這樣就摒棄了卷積神經網絡參數中浮點數所占用的內存空間[16]。學者為徹底解決浮點運算的需求提出了一種二值化網絡，將激活限制[-1，1]中，這樣就將權重和激活的網絡參數全部固定到一組固定點中，最終在ImageNet數據集上驗證了模型的有效性[17]。

2 方 法

由于醫學上實時胰腺圖像分割的需求為：在不降低或少降低分割網絡性能的前提下，盡可能減少分割網絡對硬件資源的占用，使其到達一個新的最優平衡。針對上述需求，本文提出了兩個創新：①在胰腺分割工作中，提出應用通道剪枝技術，并搭建了一個端到端框架(comprehension prune-reconstructed U-net network，CPRUN)，實驗結果表明，將模型大小壓縮至極小情況下仍保證了非常高的胰腺圖像分割精度；②為適配CPRUN模型設計了一種重構的U-net卷積神經網絡，該重構模型優點是對其進行剪枝后，在其任一卷積層中所連接的任一輸出特征通道都是對胰腺分割具有高貢獻度的前提下，充分利用到各卷積層特征通道所提供的不同分辨率特征圖信息，以保障胰腺分割精度。

本文提出一個基于通道剪枝技術與重構U-net的胰腺圖像分割的端到端框架即CPRUN模型。算法的整體流程包括對于重構的U-net卷積神經網絡進行預訓練，然后通過通道剪枝算法對重構的U-net卷積神經網絡進行模型壓縮，接著生成CPRUN模型并進行微調即再訓練，直至最終輸出模型壓縮后的CPRUN模型與胰腺分割結果。該CPRUN模型的運行流程如圖1所示。

圖1 CPRUN模型的運行流程

2.1 CPRUN模型的通道剪枝算法設計

剪枝有多種粒度不同的剪枝方法：細粒度剪枝往往是對于神經元之間的權重連接進行修剪，其進行剪枝后的模型精度往往最高；而較粗糙的粗粒度的剪枝方法，其缺點是不可避免的忽略相當數量的冗余參數。由于剪枝方法的粒度和與模型精度相互制約，為了使CPRUN模型在模型大小和分割網絡精度間取得最優平衡，決定選取通道剪枝算法來對重構的U-net卷積神經網絡進行修剪，使其一方面不會受限于細粒度剪枝所需的特定算法庫或者硬件制約；另一方面，又避免了粗粒度的剪枝算法所導致的模型壓縮效果不夠。

本文的方法是定義一個全局的閾值因子α(α∈(0,1))， 當重構的U-net卷積神經網絡預訓練到接近收斂時，保留其卷積神經網絡模型與對應模型參數。剪枝方法的第一步是，將連接在每個卷積層之后的BN(batch normalization)層的權重參數Wij集合起來組成一個集合γ，再對γ中所有元素γk取絕對值并按照大小排序，其數學表達式如式(1)所示

γ=sort({|Wij|})

(1)

其中， |Wij| 表示重構的U-net卷積神經網絡中第i層卷積層的第j條輸出特征通道所對應的BN層的權重參數再取絕對值， sort() 表示排序函數，所有排序過的 |Wij| 共同構成集合γ；第二步是統計集合中元素γk的個數numberBN， 將其與α相乘后得到集合剪切位置下標γthreshold， 其數學表達式如式(2)所示

γthreshold=α×numberBN

(2)

第三步是對集合γ中的全部元素γk進行處理，當k≥γthreshold時，在集合γ中保留元素γk。當k<γthreshold時，在集合γ中刪除元素γk。 其數學表達式如式(3)所示

(3)

最后，對那些刪除的集合元素γk，便要將其所代表的Wij所對應的卷積層輸出特征通道包括其前后輸入輸出連接全部刪掉。而集合γ中剩余的元素便代表著與其對應相連的卷積層輸出特征通道保留滿足α閾值因子的要求。最終重構的U-net卷積神經網絡轉化為CPRUN模型。

為將剪枝算法設計解釋清楚，CPRUN模型的剪枝算法的流程見表1。

表1 通道剪枝算法流程

為直觀表示通道剪枝算法的實際操作，CPRUN模型的通道剪枝算法設計如圖2所示：其中左邊虛線方框表示的是重構的U-net卷積神經網絡中第i層卷積層中經過剪枝處理后需要被減掉的特征通道，左邊實線方框表示剪枝后需要保存下來的卷積層特征通道；Wij表示的是BN層權重,其中虛線橢圓指的是經閾值因子α作用后需要被剪掉的BN層權重，相應的實線橢圓指的是保留下來的權重；Ci1至Cin表示的是第i個卷積層中的n條輸出特征通道；圖2中左邊虛線和實線方框共同構成重構的U-net卷積神經網絡的第i層卷積層，右邊實線方框共同構成第j(j=i+1) 層卷積層。

圖2 剪枝過程

BN層的作用是可以較好解決梯度消失和梯度爆炸問題，CNN(convolutional neural network)往往是在一個卷積層之后加入BN層，同時由前人實驗指出，由于BN層的特性，BN層可能是最適用于考慮進行通道剪枝的區域。所以CPRUN模型決定在BN層區域完成通道剪枝操作[18]。在通道剪枝完成之后，經過模型壓縮的卷積神經網絡模型往往會有一定程度的精度衰減與性能下降，但是通過一定數量epoch的再訓練與微調，新的剪枝模型往往可以達到甚至超越原有未壓縮模型的精度。所以還需對剪枝后生成的CPRUN模型進行微調與再訓練直至收斂并輸出分割結果。

2.2 針對通道剪枝算法的重構的U-net卷積神經網絡設計

傳統U-net分割模型的設計亮點是：通過一種跳躍連接，將分割模型下采樣環節中每兩個卷積層后輸出的特征圖信息，連接至上采樣環節對應位置的同層信息處，這樣的好處是可以利用不同分辨率的同層位置的對應特征圖信息來對目標的分割提供幫助。

CPRUN模型的設計目標指明，需要在保障模型壓縮程度的情況下，盡可能穩定或略有提高分割網絡的精度。針對上述目標，本文提出了一種重構的U-net卷積神經網絡。從2.1節可知，由于CPRUN模型考慮將所有卷積層的特征通道按重要性從大到小排序，故最終所確定的剪枝方案對于每一個卷積層的每一個輸出特征通道都有明確的刪除或保留要求。這樣經過通道剪枝算法進行模型壓縮后，分割網絡的所有卷積層的輸出特征通道都是對于胰腺分割貢獻度較高的。

至此，重構的U-net卷積神經網絡的設計方法為：①為適配通道剪枝技術，重構的U-net卷積神經網絡中卷積層的定義方法是：逐個卷積層指定其輸入特征通道數和輸出特征通道數；②將重構的U-net卷積神經網絡中下采樣環節中每個卷積層后輸出的特征圖信息，都連接至上采樣環節中對應位置的同層信息處。

由圖3展示為CPRUN模型適配設計的重構U-net結構框架：其中Conv(3×3)，ReLU表示先進行2D卷積處理，后用ReLU激活函數進行激活。MaxPool2d表示2D池化層，sigmoid表示2D激活函數。其中copy and crop 表示對特征圖的復制和裁剪。帶CPRUN copy and crop指示的加粗箭頭表示重構的U-net卷積神經網絡加入的跳躍連接，帶copy and crop指示的箭頭則表示傳統U-net分割模型的跳躍連接，帶Upsample指示的箭頭表示上采樣。從圖3不難看出，重構的U-net卷積神經網絡的跳躍連接數目是傳統U-net分割模型的兩倍，更強程度上對重構U-net同層處特征圖信息加以充分利用與提取。

2.3 評價參數

Dice系數是醫學圖像處理領域使用較為廣泛的評價指標，它可以非常好體現出預測值和標簽之間的相關關系，其定義表示如式(4)所示，這里P指的預測的分割結果，T指的是胰腺的標簽

(4)

實驗的評價參數主要是參考Dice，但是同時為了驗證實驗結果的有效性，我們同時采用了精確率Precision和感興趣區域交并比Iou評價參數進行佐證。精確率Precision的定義如式(5)所示，其中TP(true positive)表示模型預測結果為正例，實際結果為正例；FP(false positive)表示模型預測結果為正例，實際結果為負例

(5)

感興趣區域交并比Iou的定義如式(6)所示，其中P表示預測結果的集合，T表示真值的集合

(6)

為了表明壓縮后的模型與未壓縮的模型之間的壓縮程度，定義一個評價指標即壓縮率ratio， 其定義如式(7)所示

(7)

3 實 驗

3.1 實驗環境與訓練過程

實驗在Pytorch上驗證了CPRUN模型的有效性，模型的預訓練epoch次數定為30，微調epoch次數定為60。實驗環境是CPU為I7-7800X；顯卡為1080TI；64 G內存；Pytorch版本為1.5.1；CUDA版本為10.1；顯卡驅動版本為430.50。

圖3 為CPRUN模型適配設計的重構U-net模型

實驗中模型的訓練過程包括重構U-net的預訓練與后續CPRUN模型的微調，預訓練與微調過程中分割模型的損失函數設計均為-Dice， 其訓練過程中的超參數設置見表2。

模型的訓練過程中，重構U-net的精度收斂曲線與損失收斂曲線如圖4所示，CPRUN模型微調的訓練過程曲線如圖5所示，其中圖4和圖5中訓練精度表示訓練過程中模型的精度，訓練損失函數表示訓練過程中模型的損失函數，訓練次數指的是對訓練集所有數據訓練的次數。

表2 訓練過程的超參數設置

圖4 重構U-net預訓練過程曲線

圖5 CPRUN模型微調的訓練過程曲線

3.2 重構U-net的實驗設計

為了觀察重構的U-net分割網絡與傳統的U-net分割網絡之間的模型大小和性能之間的差距，設計了這兩種模型的對照實驗，其結果見表3。通過觀察可以發現，重構的U-net分割網絡優點是Dice大于傳統的U-net卷積神經網絡約一個百分點，缺點是重構的U-net分割網絡模型大小略大于傳統的U-net分割網絡。雖然重構的U-net分割網絡模型大小略有提升，但是通過后續應用剪枝技術，其模型大小會被壓縮；同時2.2節闡述了重構的U-net分割網絡是為了適配CPRUN模型的剪枝策略設計的，即增加的跳躍連接是為了充分應用后續的剪枝操作所帶來的高貢獻度的特征通道提供的不同維度特征圖信息。故重構的U-net分割網絡并不是只為了提高后續CPRUN模型的分割精度而忽視模型大小的增加。

表3 重構U-net分割網絡的對照實驗

鑒于前人的模型壓縮工作中加入了L1正則化，用來幫助訓練卷積神經網絡的權重，同時對損失函數施加約束以便于挑選對標簽特征貢獻相對較大的特征通道[18]。故本文的CPRUN模型在重構的U-net分割模型進行初始訓練時，就需要考慮是否為了后續通道剪枝的應用來加入正則化技術。正則化分為兩種即L1正則化與L2正則化(也可稱為L1范數與L2范數)，L1正則化與L2正則化可以看作是加在損失函數后的懲罰項，其定義分別如式(8)和式(9)所示

(8)

(9)

但有以下兩方面原因促使CPRUN模型在重構的U-net分割模型進行預訓練時，并不采取正則化技術進行預訓練：①CPRUN模型設計的目標中一部分的需求是盡量不給硬件增加額外的計算負擔，模型的設計應該盡量精簡；②CPRUN模型的實驗結果顯示在重構的U-net模型進行預訓練時不論加入哪種正則化技術，對于分割精度的提高都沒有任何增益，反而會拉低一些分割精度。

在實驗階段，針對于重構的U-net卷積神經網絡，設計了應用不同正則化技術的對照實驗，其結果見表4。觀察表4可知，對重構的U-net卷積神經網絡不采用L1或L2正則化，并不會影響其分割網絡在NIH胰腺數據集的性能。所以最終CPRUN模型決定不采用任何正則化技術來對重構的U-net分割網絡進行約束。

表4 正則化對重構U-net分割網絡的影響

3.3 CPRUN模型的實驗設計

為觀察CPRUN模型在不同的閾值因子α下的性能表現，對CPRUN模型設計了應用不同α的對照實驗。其結果見表5。觀察表5可知：①CPRUN模型隨著閾值因子α逐漸增大直至0.7，可以看出剪枝后模型大小也在逐步減小，同時Dice評價指標基本不變；②CPRUN模型在閾值因子α設定為0.7時，不但剪枝前后模型的Dice基本沒有變化，并且模型大小壓縮到比較小只有1.72 Mb，且壓縮率ratio達到8.0%；③當閾值因子α增大至超過0.7時，雖然模型大小進一步壓縮，但是，Dice評價指標由于此時的CPRUN模型卷積層特征通道數太少，對于胰腺特征的提取已經不夠準確且開始出現了下降，但是下降幅度很小。

同時繼續觀察表5的剩余評價指標可以得出：①隨著閾值因子α逐漸增大直至0.7，可以看出剪枝后模型大小也在逐步減小，同時Precision和Iou評價指標基本不變甚至略有增大；②當閾值因子α增大至超過0.7時，Precision和Iou評價指標也是呈現下降趨勢，與Dice評價指標的變化基本保持一致。

表5 不同壓縮程度的CPRUN模型表現

表6則詳細闡述了當閾值因子α采取不同取值時，CPRUN網絡的每層卷積層確切剩余的特征通道數。Ci表示圖3虛線框中CPRUN模型的每個卷積層輸出的特征通道數。需要指出的是表6中當α為0時表示CPRUN模型并未進行壓縮(即重構的U-net卷積神經網絡)，此時的Ci表示卷積神經網絡中每個卷積層的初始輸出特征通道數。

表6 CPRUN模型的各卷積層剩余的特征通道數

CPRUN模型對于胰腺的分割效果如圖6所示。

圖6 分割效果

3.4 CPRUN模型與其它算法對比

表7展示了CPRUN模型在NIH數據集上與另外兩種分割模型的結果比較，表7中(A6.0W0.4)表示采用量化進行模型壓縮后，原始U-net分割網絡的結構不變，只是在權重與激活參數的位數表示上進行約束。其中A6.0的意思是激活采用6個整數位表示，放棄小數位權重參數；W0.4代表權重采用4個浮點小數位表示，放棄整數位權重參數。

另外需要指出，本文實驗階段重現了fixed-point quantization(定點量化)中采用(A6.0W0.4)策略時的分割模型，其Dice大小為79.09%。另TernaryNet模型的權重參數大小應小于其實際模型大小[20]。

從表7我們可以得到如下結論：①從表7可以觀察到當CPRUN模型的Dice和fixed-point quantization中采用(A6.0W0.4)策略時Dice大小沒有明顯差距時，可以明顯看到我們的模型大小(841 Kb)遠小于其模型大小(2.31 Mb)，約為fixed-point quantization模型大小的1/3大小；②當CPRUN模型的大小(841 Kb)和TernaryNet的權重參數大小(0.66 Mb)差距不大時，可以看到CPRUN模型的Dice比其高出9個百分點左右；③當CPRUN模型的大小壓縮至208 Kb時,可以看到CPRUN模型的Dice大小為78.31%，仍維持很高的分割水平，說明我們的CPRUN模型在模型大小和分割精度之間取得了新的最優平衡，真正做到了模型大小與分割精度共同兼顧。

表7 方法效果對比

4 結束語

本文提出了應用于醫學圖像分割中胰腺器官領域的CPRUN模型并在NIH胰腺數據集上做了驗證。CPRUN模型一方面使用了重構的U-net卷積神經網絡，另一方面集成了通道剪枝策略，從而實現了一個基于重構的U-net卷積神經網絡與通道剪枝技術的胰腺分割端到端框架。實驗結果表明，CPRUN模型可以在模型大小壓縮至極小時保持較高的分割精度，真正在胰腺圖像的分割模型大小與分割精度之前取得了新的最優平衡。進一步的工作將會在以下兩方面進行：①未來會考慮在剪枝的基礎上，融合模型壓縮的量化技術對CPRUN模型進一步進行壓縮；②同時也會關注于在時間壓縮方面有所進步，真正的為推動實時醫療夯實基礎。