基于改進(jìn)U-Net網(wǎng)絡(luò)的神經(jīng)元分割算法

程維東，葉 曦*，王 芳，平晶晶，錢同惠，張志瑋

（江漢大學(xué) a.智能制造學(xué)院；b.人工智能學(xué)院；c.設(shè)計(jì)學(xué)院，湖北 武漢 430056）

0 引言

圖像分割在生物醫(yī)學(xué)圖像分析中有著重要作用，也是醫(yī)學(xué)圖像處理方法的重要組成之一。隨著計(jì)算機(jī)硬件性能的提高，基于U-Net［1］的圖像分割技術(shù)已經(jīng)成為醫(yī)學(xué)圖像處理的有力工具，被廣泛應(yīng)用于各種醫(yī)學(xué)圖像分割任務(wù)中。目前，U-Net對視網(wǎng)膜、腦組織、肺部和血流圖像等都有較好的分割效果，也在生物圖像分析中起到了重要的作用。

神經(jīng)元結(jié)構(gòu)分割是目前生物醫(yī)學(xué)圖像分析的重要課題之一［2］。超高分辨率的電鏡圖像越來越多，但對電鏡圖像進(jìn)行有效分割依然存在挑戰(zhàn)。采用人工分割的方式雖然能獲得很好的分割效果，但這一過程十分繁瑣且耗時甚高［3］。2015年開始，U-Net在醫(yī)學(xué)圖像分割中的應(yīng)用取得了突破性進(jìn)展，從此獲得了廣泛的關(guān)注和應(yīng)用，進(jìn)而基于深度學(xué)習(xí)的分割算法成為醫(yī)學(xué)圖像分割算法的主流方法［1，4］。因此，建立合適的深度學(xué)習(xí)模型對圖像分割進(jìn)行精確、快速、有效的神經(jīng)元結(jié)構(gòu)分割，對于生物醫(yī)學(xué)圖像分析有著重大意義。

深度學(xué)習(xí)圖像分割模型從2012年開始，Hinton和學(xué)生Krizhevsky設(shè)計(jì)的AlexNet在Image Net比賽上取得了突破性的識別準(zhǔn)確率后［5］，卷積神經(jīng)網(wǎng)絡(luò)（convolutional neural network，CNN）開始成為圖像處理的主要工具。2015年，網(wǎng)絡(luò)結(jié)構(gòu)經(jīng)歷了從AlexNet到Oxford Visual Geometry Group提出的VGG模型［6］的轉(zhuǎn)變。該模型的亮點(diǎn)在于用3*3、1*1的卷積核替代了5*5卷積，在相同的感受野下，3*3的卷積參數(shù)更少。之后，U-Net在MICCAI 2016被擴(kuò)展為3D U-Net應(yīng)用于三維圖像分割［7］，該深度學(xué)習(xí)網(wǎng)絡(luò)在2018年仍然能在Multimodal Brain Tumor Segmentation Challenge 2018中取得第二名的好成績。Top2的網(wǎng)絡(luò)模型［8］在U-Net的基礎(chǔ)上做出了一些改動，該模型與U-Net同樣采用4次下采樣，通過Dice loss可以提升模型的Dice的得分，并使用Leaky Re LU激活函數(shù)提高性能。因此，在醫(yī)學(xué)圖像分割中，損失函數(shù)的設(shè)計(jì)與選用尤為重要［9］，其中用得較多的損失函數(shù)為Dice損失函數(shù)（Dice coefficient loss）［10］、交叉熵?fù)p失函數(shù)（cross entropy loss）、相對熵?fù)p失函數(shù)（KL loss）等。目前，一些基于CNN的算法對神經(jīng)元重要細(xì)節(jié)分割上有所缺失，分割圖像中像素間的關(guān)聯(lián)度不強(qiáng)，導(dǎo)致線條也存在斷裂。因此，如何利用注意力機(jī)制，對電鏡神經(jīng)元圖像中神經(jīng)元的某些特征進(jìn)行重點(diǎn)學(xué)習(xí)、怎樣針對神經(jīng)元圖像選用合適的損失函數(shù)或者組合損失函數(shù)，都是值得研究的重要問題。

綜上所述，本文提出一種用于電鏡圖像中神經(jīng)元結(jié)構(gòu)分割的深度學(xué)習(xí)模型：Attention-KL U-Net（A-KLU-Net）。該模型主要由4個下采樣、4個上采樣、一個自注意力模塊組成，并使用組合損失函數(shù)來評估模型。本文以含有BN層和卷積時進(jìn)行邊緣補(bǔ)零的U-Net結(jié)構(gòu)為骨架，經(jīng)過大量實(shí)驗(yàn)發(fā)現(xiàn)自注意力機(jī)制的使用可以有效地避免神經(jīng)元分割結(jié)果圖中的斷裂和多余部分的出現(xiàn)，并進(jìn)一步提出了改進(jìn)的注意力機(jī)制模型來加強(qiáng)對細(xì)節(jié)的學(xué)習(xí)；其次，在網(wǎng)絡(luò)模型中使用了Dice與KL散度疊加的損失函數(shù)，使得模型可以更好地學(xué)習(xí)。

1 模型建立

1.1 U-Net模型

U-Net網(wǎng)絡(luò)的核心思想是利用較少的訓(xùn)練集進(jìn)行端對端的訓(xùn)練［11］，U-Net模型是建立在全連接卷積神經(jīng)網(wǎng)絡(luò)（fully connected network，F(xiàn)CN）［12］架構(gòu)上，對FCN結(jié)構(gòu)進(jìn)行修改和更為精細(xì)設(shè)計(jì)。但該網(wǎng)絡(luò)不能很好地恢復(fù)經(jīng)池化層之后的圖像細(xì)節(jié)結(jié)構(gòu)信息，易造成分割斷裂的現(xiàn)象［11］?；诖?，本研究將在U-Net模型的基礎(chǔ)上添加邊緣補(bǔ)零和BN（batch normalization）層，通過在卷積過程中對圖像進(jìn)行邊緣補(bǔ)零操作，使得網(wǎng)絡(luò)更好地考慮到圖像的邊緣像素點(diǎn)；同時，圖像的尺寸在卷積后保持原來大小，只在用最大池化層對圖像特征進(jìn)行池化時，改變圖像分辨率。添加BN層主要是將每個通道輸出數(shù)據(jù)的均值和方差進(jìn)行歸一化處理，使得其均值接近0、方差接近1，從而提升關(guān)鍵性指標(biāo)F1，更好地提升網(wǎng)絡(luò)分割性能。

1.2 采樣模型

模型中下采樣和上采樣過程中每層的卷積結(jié)構(gòu)都含有卷積層（Conv）、BN層以及線性整流單元（rectified linear unit，Re LU）激活層。與原始U-Net不同，本文模型每次卷積會添加zero padding，通過在卷積過程中對圖像進(jìn)行邊緣補(bǔ)零操作，使得圖像的尺寸在采樣過程中不變，同時進(jìn)一步讓網(wǎng)絡(luò)考慮邊緣像素增強(qiáng)對細(xì)節(jié)的敏感度。下采樣對圖像特征進(jìn)行提取，上采樣對圖像進(jìn)行還原，用轉(zhuǎn)置卷積（transpose convolution，TSConv）進(jìn)行上采樣，再與下采樣對應(yīng)的特征圖拼接，對其卷積縮小通道數(shù)。模型如圖1所示（圖中卷積核大小為3*3）。

圖1 上采樣和下采樣模型Fig.1 Upsampling and downsampling models

線性整流單元ReLU為

1.3 注意力機(jī)制及其改進(jìn)模型

注意力機(jī)制選擇性地關(guān)注整體信息的一部分，獲得需要重點(diǎn)關(guān)注的區(qū)域，有效地提高了人類識別和認(rèn)知的效率［13-15］。同理，在深度學(xué)習(xí)中提取到的信息流以同等重要性向后流動，通過某些先驗(yàn)信息，抑制某些無效信息的流動，使重要信息保留。因此，本文選取自注意力（self-attention）機(jī)制來對圖像中的特征進(jìn)行針對性學(xué)習(xí)，并對注意力模塊進(jìn)行了改進(jìn)。模型如圖2（a）所示，圖中注意力機(jī)制由卷積層（Conv2d：卷積核尺寸為1*1卷積核）、BN層、激活層（激活函數(shù)為Sigmoid函數(shù)）組成。通過卷積降低特征圖的通道數(shù)，使得特征維度為1；再使用BN層做歸一化；最后，通過Sigmoid函數(shù)來進(jìn)行激活，由于Sigmoid函數(shù)值域在（0，1）使得數(shù)據(jù)波動小，會生成一個合適的模板。

Sigmoid激活函數(shù)為

本文還通過實(shí)驗(yàn)發(fā)現(xiàn)在注意力機(jī)制中，把原圖與最后一次上采樣的結(jié)果進(jìn)行拼接，用拼接的特征圖卷積得到單通道的特征圖，再與最后一次上采樣的輸出融合，可以很好地提高分割圖像中需要的細(xì)節(jié)，像素點(diǎn)關(guān)聯(lián)度也更好，故本文將使用改進(jìn)后的自注意力模型。具體結(jié)構(gòu)如圖2（b）所示：將原圖在第二維度與其拼接，給網(wǎng)絡(luò)提供更多的圖像細(xì)節(jié)和像素間的關(guān)聯(lián)度，該方法理論上具有一定的通用性。

圖2 改進(jìn)前和改進(jìn)后注意力模型Fig.2 Attention model before and after improvement

1.4 損失函數(shù)的組合

損失函數(shù)是機(jī)器學(xué)習(xí)里最為基礎(chǔ)和關(guān)鍵的要素，決定了反向傳播更新各層神經(jīng)元連接之間的權(quán)重和偏差。本文結(jié)合訓(xùn)練數(shù)據(jù)，給出一個較優(yōu)的損失函數(shù)。Dice本質(zhì)上屬于二分類，用于衡量兩個樣本的重疊部分，該指標(biāo)屬于［0，1］，越接近1重疊度越高。KL散度用于度量兩個分布之間相似性。網(wǎng)絡(luò)在選取損失函數(shù)時需要考慮數(shù)據(jù)集圖像的類別、細(xì)致度、分割目標(biāo)大小等多方面因素。本文分別對3種損失函數(shù)進(jìn)行了實(shí)驗(yàn)，選取最為契合的一種。實(shí)驗(yàn)研究發(fā)現(xiàn)損失函數(shù)采用Dice loss和KL loss的組合能提高F1指標(biāo)，從而提升網(wǎng)絡(luò)性能。

1.5 A-KLU-Net模型

本文以U-Net為骨架與本研究提出的改進(jìn)注意力模型結(jié)合構(gòu)建了A-KLU-Net模型。如圖3所示，該模型主要由4個下采樣（池化層的卷積核為2*2）、4個上采樣、一個Self-Attention模塊組成。損失函數(shù)采用Dice loss與KL loss疊加的形式。圖3中Conv2d為卷積層，Maxpool2d為最大池化層。

圖3 A-KLU-Net模型示意圖Fig.3 Schematic diagram of A-KLU-Net model

1.6 算法框架

算法框架主要分兩大部分：網(wǎng)絡(luò)訓(xùn)練和測試（見圖4）。網(wǎng)絡(luò)訓(xùn)練部分首先將圖像進(jìn)行幾何變換（旋轉(zhuǎn)和翻轉(zhuǎn)），在不改變原始圖像的條件下進(jìn)行了數(shù)據(jù)增廣，提升網(wǎng)絡(luò)模型的魯棒性。把訓(xùn)練集輸入A-KLU-Net模型，對網(wǎng)絡(luò)的輸出用Softmax函數(shù)對圖像的每個像素進(jìn)行分類得到背景和細(xì)胞膜兩種圖，之后用損失函數(shù)對預(yù)測圖和標(biāo)簽進(jìn)行評估的同時通過Validation模塊測試網(wǎng)絡(luò)得分，通過Adam優(yōu)化器優(yōu)化，直到（5）式中l(wèi) o s s值降到最低，模型訓(xùn)練完成，保存在Validation模塊中得分最高的網(wǎng)絡(luò)。在訓(xùn)練過程中，為了提高效率需要給網(wǎng)絡(luò)設(shè)置一個耐性值patience，即在patience個epoch網(wǎng)絡(luò)性能沒有提高就結(jié)束訓(xùn)練。測試過程：將測試集的電鏡神經(jīng)元圖像輸入保存的訓(xùn)練效果最佳網(wǎng)絡(luò)，對圖像進(jìn)行預(yù)測，從而達(dá)到分割的效果。

圖4 網(wǎng)絡(luò)結(jié)構(gòu)框圖Fig.4 Network structure block diagram

2 實(shí)驗(yàn)結(jié)果與分析

2.1 實(shí)驗(yàn)數(shù)據(jù)與基本參數(shù)

本研究采用的數(shù)據(jù)集為ISBI 2012提供的數(shù)據(jù)集［16-17］，訓(xùn)練數(shù)據(jù)是來自果蠅第一齡幼蟲腹側(cè)神經(jīng)索（VNC）的連續(xù)部分透射電子顯微鏡（ssTEM）數(shù)據(jù)集的30個部分，共30張圖片，每張圖片像素大小為512*512，標(biāo)簽為二值圖像，即分割對象的像素為黑色（細(xì)胞膜），其余像素為白色。本文采用旋轉(zhuǎn)和翻轉(zhuǎn)對數(shù)據(jù)集進(jìn)行增廣以減小過擬合的發(fā)生。網(wǎng)絡(luò)模型實(shí)現(xiàn)的軟件和模塊分別是Python3.7.10、Torch1.8.1和CUDA10.0，并且使用NVIDIA GeForce RTX 3090 GPU進(jìn)行訓(xùn)練。網(wǎng)絡(luò)基本參數(shù)：優(yōu)化器為Adam，學(xué)習(xí)率為0.001，訓(xùn)練Batch為4，測試Batch為9。

2.2 模型評估指標(biāo)

本文為了更好地評估網(wǎng)絡(luò)的分割效果，采用以下4個評估指標(biāo)：Accuracy（A）表示正確性，即分類背景和細(xì)胞膜的像素在整個圖像像素中的百分比；Precision（P）表示精確度，即測試樣本中準(zhǔn)確預(yù)測為細(xì)胞膜的數(shù)除以準(zhǔn)確預(yù)測的總數(shù)；Recall（R）表示召回率又稱為敏感度；醫(yī)學(xué)圖像分割常用指標(biāo)Dice通用式：F1表示預(yù)測圖與真實(shí)圖的重疊率。

式中，T P表示預(yù)測為細(xì)胞膜，標(biāo)簽為細(xì)胞膜；T N表示預(yù)測為背景，標(biāo)簽為背景；F P表示預(yù)測為細(xì)胞膜，標(biāo)簽為背景；F N表示預(yù)測為背景，標(biāo)簽為細(xì)胞膜。

2.3 結(jié)果分析

2.3.1 自注意力機(jī)制與損失函數(shù)對模型的影響 為了進(jìn)一步分析本文算法中的自注意力機(jī)制與損失函數(shù)對模型的影響，本研究通過對損失函數(shù)和自注意力機(jī)制的組合、不同損失函數(shù)的組合分析可知：①KL散度可以有效地提升正確性，自注意力機(jī)制可以有效提高正確性和精確度。②如表1所示，前三行數(shù)據(jù)與后三行都對不同的損失函數(shù)進(jìn)行了比較，發(fā)現(xiàn)單獨(dú)使用Dice損失函數(shù)比交叉熵?fù)p失函數(shù)（Ce loss）準(zhǔn)確度上低，但Dice可以有效地提高F1指標(biāo)和召回率，且與自注意力機(jī)制模型結(jié)合時，在關(guān)鍵性指標(biāo)F1上也有一定的提升?；诖?，本文采用組合損失函數(shù)方式，把Dice和交叉熵與KL散度進(jìn)行適當(dāng)?shù)寞B加，通過實(shí)驗(yàn)研究發(fā)現(xiàn)，當(dāng)Dice與0.1倍的KL散度結(jié)合時，分割效果較之單獨(dú)的Dice有所提升，具體對比表1的第二行和第五行，發(fā)現(xiàn)其對A、F1指標(biāo)有一定提升。為提高效率，本文添加了Validation模塊，在20代內(nèi)網(wǎng)絡(luò)無法提升就結(jié)束訓(xùn)練。為了避免測試集和驗(yàn)證集可能對實(shí)驗(yàn)的影響，采用不同的兩種數(shù)據(jù)集比例進(jìn)行實(shí)驗(yàn)。具體數(shù)據(jù)見表1（n=0.1）。

由表1可知，表中第四行和第五行帶有自注意力機(jī)制，使得模型對圖像細(xì)節(jié)分割的指標(biāo)P有所提升，從而很好地提升了本文關(guān)鍵性指標(biāo)F1。圖5（a）為經(jīng)過self-attention模塊之后的效果圖，圖5（b）和圖5（c）從左到右分別是原始圖、標(biāo)準(zhǔn)圖、預(yù)測圖。通過對比圖5（a）、圖5（b）和圖5（c）中紅色方框的分割情況，可以看出有self-attention模塊對圖像的復(fù)雜區(qū)域分割的效果更好，給后續(xù)像素分類提供了更多的細(xì)節(jié)；上述差別反映到網(wǎng)絡(luò)模型評價指標(biāo)中，可以發(fā)現(xiàn)圖像的像素分類精確度（P）有很好的提升。所以，在最后一次上采樣結(jié)束后，使用Self-Attention可以有效減少圖像信息的丟失，在恢復(fù)圖像分辨率時更好地分割出圖像的細(xì)節(jié)，從而提高網(wǎng)絡(luò)模型的分割效果。

表1 自注意力與損失函數(shù)效果對比Tab.1 Effect comparison of self-attention and loss function

圖5 Self-attention效果對比圖Fig.5 Self-attention effect comparison graph

2.3.2 實(shí)驗(yàn)分析 為了與現(xiàn)有算法U-Net、文獻(xiàn)［11］、文獻(xiàn)［18］比較，實(shí)驗(yàn)采用25張圖作為訓(xùn)練集，5張圖作為測試集。如表2所示，本文和現(xiàn)有算法在分割效果上較之原始U-Net都有較好的提高。同時，本文與文獻(xiàn)［11］相比，在正確性（A）、召回率（R）上分別提高了0.018 33和0.003 04，并在其基礎(chǔ)上添加了精確度（P）對分割的精確度進(jìn)行評估，相對U-Net提高了0.013 99；相對于文獻(xiàn)［18］的方法用于本數(shù)據(jù)集，在精確度上稍有不足，但在其他3個指標(biāo)上都有明顯優(yōu)勢。文獻(xiàn)［11］將分組卷積與改進(jìn)的U-Net網(wǎng)絡(luò)結(jié)構(gòu)相結(jié)合，其F1指標(biāo)與本文算法基本持平，但本文算法在準(zhǔn)確度和召回率上都高于該算法。映射到分割圖像上，即圖像分割得更為準(zhǔn)確和細(xì)致。

表2 不同方法比較Tab.2 Comparison of different methods

圖6為本文算法與原始U-Net在ISBI 2012數(shù)據(jù)集中的部分圖像的分割效果圖。通過對圖像彩色化可使圖像可視效果更好。其中，圖6從左到右分別是原始圖像、標(biāo)準(zhǔn)圖和3種方法的預(yù)測圖。通過觀察并將圖像彩色化可以看出本文模型在細(xì)胞膜分割細(xì)節(jié)方面更好且大部分地方無明顯斷裂現(xiàn)象。通過對比可以看出本文在細(xì)胞膜分割的粗細(xì)上，更加與標(biāo)準(zhǔn)圖匹配，在細(xì)節(jié)上也更加匹配且與label對比多余部分更少。本文算法更好地避免了細(xì)胞膜分割的斷裂和多余部分的出現(xiàn)，具體如圖（6）中的紅色標(biāo)記點(diǎn)所示。

圖6 分割結(jié)果圖Fig.6 Segmentation result graph

3 結(jié)語

本文提出了對神經(jīng)元圖像的深度學(xué)習(xí)分割模型A-KLU-Net，該網(wǎng)絡(luò)在ISBI 2012數(shù)據(jù)集上進(jìn)行了實(shí)驗(yàn)，其正確率、F1指標(biāo)、精確度和召回率分別達(dá)到0.930 43、0.956 79、0.953 26、0.960 34。通過損失函數(shù)和自注意力機(jī)制的組合、不同損失函數(shù)的組合分析，驗(yàn)證了改進(jìn)的自注意力機(jī)制可以有效提升圖像分割的準(zhǔn)確度，Dice與KL散度的結(jié)合對本研究的數(shù)據(jù)集的效果更好。后續(xù)可以考慮通過加入多尺度輸入輸出改善網(wǎng)絡(luò)高層特征逐步丟失初始信息的問題，并且多個輸出也有利于網(wǎng)絡(luò)的梯度反向傳播、改善訓(xùn)練效率，也可以在多尺度特征融合。