LKA_Unet：將基于大卷積核的空間注意力機制應用于Unet

蔡一民 楊威

摘? 要： 醫(yī)學圖像分割技術能夠自動勾勒出醫(yī)學圖像中人體組織結(jié)構和病變區(qū)域，進而輔助醫(yī)生診療，減少誤診率。然而大多數(shù)現(xiàn)有醫(yī)學圖像分割模型使用的是感受野固定且很小的卷積核，不能完美地擬合圖像中的全局特征信息。針對這個問題，我們提出了一個具有超大感受野的空間注意力機制模塊。在DentalPanoramicXrays數(shù)據(jù)集上的實驗表明，該模塊加入到Unet網(wǎng)絡后，分割精度提升了約2個百分點。

關鍵詞： 深度學習； 醫(yī)學圖像分割； 空洞卷積； 空間注意力機制

中圖分類號：TP391.1? ? ? ? ? 文獻標識碼：A? ? ?文章編號：1006-8228（2023）05-81-04

LKA_Unet： applying spatial attention mechanism based on large

convolutional kernel to Unet

Cai Yimin， Yang Wei

（Guizhou University， Guiyang， Guizhou 550025， China）

Abstract： Medical image segmentation technology can automatically outline the human tissue structure and lesion areas in medical images， thus assisting doctors in diagnosis and treatment and reduce misdiagnosis. However， most existing medical image segmentation models use convolutional kernels with fixed and small receptive fields， which cannot perfectly fit the global feature information in images. To address this problem， we propose a spatial attention mechanism module with a large field of perception. Experiments on the DentalPanoramicXrays dataset show that the segmentation accuracy is improved by about 2 percentage points after the module is added to the Unet network.

Key words： deep learning; medical image segmentation; atrous convolution; spatial attention mechanism

0 引言

圖像分割技術是人工智能中一個重要分支。醫(yī)學圖像分割在腫瘤、心血管等領域應用廣泛。然而由于醫(yī)學圖像結(jié)構復雜且存在大量非結(jié)構化信息（如圖像陰影），醫(yī)學圖像分割任務往往是非常困難的。現(xiàn)有的醫(yī)學圖像分割模型考慮到模型參數(shù)量以及訓練速度等因素，卷積核的大小通常都是[3×3]、[5×5]、[7×7]、[9×9]等。這些模型并不能完美地擬合學習圖像中的長距離依賴信息。受VAN[1]的啟發(fā)，本文提出一個基于超大卷積核的空間注意力機制模塊，并探索了該模塊對經(jīng)典Unet模型分割精度的提升程度。

1 相關工作

深度學習是人工智能領域的一個重要分支，利用大量數(shù)據(jù)進行深度學習以獲得較高的分類、分割和目標檢測的精度。深度學習的出現(xiàn)改變了人工智能的發(fā)展歷程，其優(yōu)勢在于，能快速準確地訓練出強大的模型，其本質(zhì)是通過對大量數(shù)據(jù)進行訓練。在經(jīng)過近幾年的發(fā)展后，深度學習已經(jīng)從圖像分類領域延伸到了圖像分割、文本識別和語音語義識別等領域。在過去十年中，大量科研工作者投入到了深度學習技術研究工作中，包括最初的AlexNet網(wǎng)絡[2]，以及后續(xù)的Resnet網(wǎng)絡[3]、VIT網(wǎng)絡[4]和Swin Transformer[5]網(wǎng)絡等。深度學習中廣泛使用的方法是卷積神經(jīng)網(wǎng)絡（Convolutional Neural Network， CNN），其基本思想是利用多個卷積和不同大小的池化層將圖像表示為一個個獨立且相似度較高的值，將這些值合并后生成最終輸出。通過多次卷積運算產(chǎn)生網(wǎng)絡訓練樣本從而提高模型性能。CNN算法簡單易用，且能在較短時間內(nèi)得到高精度圖像分割結(jié)果。

在各種醫(yī)學圖像分割任務上，U-Shape架構的網(wǎng)絡，也被稱為類U-Net[6]的網(wǎng)絡，已經(jīng)取得了較大的成功并成為分割任務的標桿。在Unet的基礎上，衍生出了各種魔改模型，包括但不限于V-Net[7]，Unet++[8]，Unet3+[9]，Swin Unet[10]，Res Unet[11]等。還有人提出了R-shape架構，該類模型包括UnetR[12]，Swin UnetR[13]和UNetR++[14]等。上述這些模型在不同分割任務上取得了巨大的成功，并且在商業(yè)應用上也取得了令人矚目的成就。但是，在這些模型內(nèi)部使用的卷積核幾乎完全是固定的。然而，在圖像分割整個過程中，圖像中的信息密度是會發(fā)生變化的。從直觀上而言，高分辨率的醫(yī)學圖像中特征信息比較分散，而低分辨率的醫(yī)學圖像中特征信息比較密集。固定大小的卷積核并不能很好地擬合密度不同的特征信息。

在本文中，我們提出了一個新的空間注意力機制模塊，該模塊被添加到經(jīng)典的Unet網(wǎng)絡中以試圖緩解上述問題。在Unet網(wǎng)絡的不同階段中這個注意力機制模塊的感受野大小是可變的。當圖像處于高空間分辨率的階段我們用大感受野的卷積核去提取特征，當圖像處于低空間分辨率的時候我們使用小感受野的卷積核去提取特征。實驗結(jié)果表明，在經(jīng)典的Unet模型中添加了我們提出的空間注意力機制模塊，分割精度提高了將近2個百分點。

2 模型設計細節(jié)

2.1 模型總覽

LKA_Unet模型架構如圖1所示，該模型由四個編碼器（Down1、Down2、Down3以及Down4）和四個解碼器（Up1、Up2、Up3、Up4）組成。每個編碼器由一個MLP子模塊、一個Attention子模塊和一個DownSample子模塊組成。每個解碼器由一個MLP模塊、一個Attention子模塊和一個UpSample子模塊組成。我們將在2.2節(jié)對每個子模塊的內(nèi)部結(jié)構以及主要功能進行說明。

2.2 子模塊介紹

2.2.1 MLP模塊

如圖2所示，按照卷積層，實例歸一化層和一個ReLU激活層的順序堆疊兩次后構成了MLP子模塊。MLP模塊并不會改變特征圖的空間分辨率，它改變的是特征圖的通道數(shù)目。不同的通道可以被視為對不同特征的響應，每個通道都有一個對應的卷積核，這些卷積核的參數(shù)是可學習的。

2.2.2 Attention模塊

Attention模塊是LKA_Unet模型的主要亮點。針對不同空間分辨率的圖像，Attention模塊的感受野的大小是變化的。具體來說，在Down1中，圖像的空間分辨率和原始圖像相同（deal_in模塊作為一個中間模塊，只會改變原始輸入圖像的通道數(shù)，以便于模型在代碼實現(xiàn)方面更為簡潔），因此它里面的Attention模塊的感受野大小為31×31。在Down2中，此時圖像被下采樣到遠空間分辨率1/4，因此在Down2中，Attention模塊的感受野大小為15×15。在Down3和Down4中，輸入圖像的空間分辨率分別相當于原始圖像的1/8和1/16，因此它倆內(nèi)部的Attention模塊感受野大小分別為7×7大小和3×3。

由于圖像經(jīng)過在不同的編碼器中，空間分辨率也不一樣。直觀上看，空間分辨率較高的圖像需要使用較大的卷積核去提取特征，而對于空間分辨率較小的圖像需要用較小的卷積核來進行特征提取。如果使用固定大小的卷積核，卷積核過小，則無法提取高分辨率圖像中的長距離依賴信息，若卷積過大，則會在提取低分辨率圖像中的特征時導致無用計算過多，甚至造成因卷積核過多零填充導致模型的輸出異常（比如輸出中含有過多的nan或者inf）。

在實現(xiàn)方面，針對較大的卷積核，我們是通過堆疊多層空洞卷積來實現(xiàn)的。對于一個堆疊了多層空洞卷積的Attention模塊，其感受野的大小可以由如下遞推公式計算出來：

[Kn'=Kn+（Kn-1）×（dn-1）][F1=K1'][Fn]

[=Fn-1+Kn'-1][=Fn-1+（Kn-1）×dn]

其中，[Kn]表示第n個的卷積核大小，[dn]表示第n個卷積核的空洞大小，[Fn]表示經(jīng)過堆疊n個卷積核之后Attention模塊感受野的大小。Attention是一個空間注意力機制模塊，因此，它不會改變輸入圖像的空間分辨率，它惟一的作用是提取圖像特征。借助Attention模塊的輸入和輸出之間的相等關系以及上述公式，就可以設計出堆疊的卷積核個數(shù)n，以及每個卷積核的大小[Kn]和空洞參數(shù)[dn]。

2.2.3 DownSample子模塊和UpSample子模塊

DownSample子模塊主要用于降低圖像空間分辨率，即下采樣。UpSample子模塊主要用于恢復圖像空間分辨率，即上采樣。DownSample由一個卷積層和一個實例歸一化層組成，UpSample由一個轉(zhuǎn)置卷積層和一個實例歸一化層組成。在代碼實現(xiàn)方面也可以通過bilinear插值的方式替代UpSample模塊里面的轉(zhuǎn)置卷積層。在Down1中和Up1中將卷積核大小k和步長s均設置為4，分別用于4倍下采樣圖像和4倍上采樣圖像。在Down2、Down3、Down4和Up4、Up3，以及Up2中卷積核大小k和步長s分別設置為2，分別用于2倍下采樣和2倍上采樣。從Down1到Down4階段，下采樣倍數(shù)總計為32倍，和經(jīng)典Unet完全一樣。從Up4到Up1，上采樣倍數(shù)總計為32，完成恢復到原圖像的空間分辨率。

3 實驗細節(jié)

3.1 實驗數(shù)據(jù)集簡介

實驗中使用的數(shù)據(jù)集是DentalPanoramicXrays[15，16]，數(shù)據(jù)下載鏈接為：https：//github.com/SerdarHelli/Segmentation-of-Teeth-in-Panoramic-X-ray-Image-Using-U-Net。該數(shù)據(jù)集一共包含116個口腔CBCT全景片，分割目標為所有牙齒區(qū)域。該數(shù)據(jù)集總共包含116個病例，為了讀取和處理方便，作者已經(jīng)將X射線圖像轉(zhuǎn)換成了png格式的圖像并給出了對應的Ground Truth。

3.2 實驗環(huán)境和實驗細節(jié)

我們的實驗室基于Monai0.9.1[17]和Pytorch-lighting1.7.5[18]完成的，實驗使用的電腦是一臺小型工作站，該機器的操作系統(tǒng)為Windows10，CPU型號為Intel（R） Xeon（R） CPU E5-2690，顯卡為英偉達RTX2080Ti。

所有圖像尺寸被縮放到[512×512]，使用了固定隨機種子42初始化Pytorch[19]環(huán)境和Cuda環(huán)境，并使用該種子劃分訓練集和驗證集。劃分比例為0.8：0.2，其中驗證集不參與訓練。數(shù)據(jù)增強方面，我們使用了Monai[17]提供的隨機水平翻轉(zhuǎn)和垂直翻轉(zhuǎn)，以及隨機強度縮放和平移。每個模型最少訓練20個epoch，并使用Pytorch-lightning[18]提供的EarlyStop技術防止模型過擬合。我們的模型和其他對比模型在驗證集上取得的最高平均dice系數(shù)如表1所示，訓練過程中各個模型在驗證集上取得的平均dice系數(shù)變化曲線如圖3所示。

3.3 實驗結(jié)果分析

各個模型在實驗使用到的驗證數(shù)據(jù)集上取得的平均dice系數(shù)如表1所示。相比較于Unet，LKA_Unet在驗證數(shù)據(jù)集上取得的dice系數(shù)提高了2個百分點左右。相比較于SegResNet，LKA_Unet的分割性能雖然提升不是很大，但是收斂速度明顯快于SegResNet。

4 結(jié)論

在本文中我們提出了一個新的空間注意力機制模塊，該模塊被添加到經(jīng)典的Unet網(wǎng)絡后能夠促進Unet在不同階段更好地提取圖像特征信息。在DentalPanoramicXrays[15，16]數(shù)據(jù)集上進行的實驗表明，我們提出的空間注意力機制模塊確實在一定程度上提高了Unet網(wǎng)絡的分割精度。

參考文獻（References）：

[1] Guo M， Lu C， Liu Z， et al. Visual attention network[J].

arXiv preprint arXiv：2202.09741，2022

[2] Krizhevsky A， Sutskever I， Hinton G E. Imagenet

classification with deep convolutional neural networks[J]. Communications of the ACM，2017，60（6）：84-90

[3] He K， Zhang X， Ren S， et al. Deep residual learning for

image recognition[C]，2016

[4] Dosovitskiy A， Beyer L， Kolesnikov A， et al. An image is

worth 16x16 words： Transformers for image recognition at scale[J]. arXiv preprint arXiv：2010.11929，2020

[5] Liu Z， Lin Y， Cao Y， et al. Swin transformer： Hierarchical

vision transformer using shifted windows[C]，2021

[6] Ronneberger O， Fischer P， Brox T. U-net： Convolutional

networks for biomedical image segmentation[C].Springer，2015

[7] Milletari F， Navab N， Ahmadi S. V-net： Fully convolutional

neural networks for volumetric medical image segmentation[C].IEEE，2016

[8] Zhou Z， Siddiquee M M R， Tajbakhsh N， et al. Unet++：

Redesigning skip connections to exploit multiscale features in image segmentation[J]. IEEE transactions on medical imaging，2019，39（6）：1856-1867

[9] Huang H， Lin L， Tong R， et al. Unet 3+： A full-scale

connected unet for medical image segmentation[C].IEEE，2020

[10] Cao H， Wang Y， Chen J， et al. Swin-unet： Unet-like

pure transformer for medical image segmentation[J]. arXiv preprint arXiv：2105.05537，2021

[11] Diakogiannis F I， Waldner F C C O， Caccetta P， et al.

ResUNet-a： A deep learning framework for semantic segmentation of remotely sensed data[J]. ISPRS Journal of Photogrammetry and Remote Sensing，2020，162：94-114

[12] Hatamizadeh A， Tang Y， Nath V， et al. Unetr：Transformers for 3d medical image segmentation[C]， 2022

[13] Hatamizadeh A， Nath V， Tang Y， et al. Swin unetr： Swin

transformers for semantic segmentation of brain tumors in mri images[C].Springer，2022

[14] Shaker A， Maaz M， Rasheed H， et al. UNETR++：

Delving into Efficient and Accurate 3D Medical Image Segmentation[Z]. arXiv，2022

[15] HELL I S， Hamamci A. Tooth Instance Segmentation on

Panoramic Dental Radiographs Using U-Nets and Morphological Processing[J]. D{＼"u}zce {＼"U}niversitesi Bilim ve Teknoloji Dergisi，2022，10（1）：39-50

[16] Abdi A， Kasaei S. Panoramic dental X-rays withsegmented mandibles[J]. Mendeley Data， v2，2020[17] Consortium M. MONAI： Medical Open Network for AI[Z].2020

[18] Falcon W， Team T P L. PyTorch Lightning[Z].2019

[19] Paszke A， Gross S， Massa F， et al. Pytorch： An

imperative style， high-performance deep learning library[J]. Advances in neural information processing systems，2019，32