基于Group-Depth U-Net的電子顯微圖像中神經(jīng)元結(jié)構(gòu)分割

李玉慧，梁創(chuàng)學(xué)，李軍

華南師范大學(xué)物理與電信工程學(xué)院，廣東廣州510006

前言

神經(jīng)元結(jié)構(gòu)分割是目前生物醫(yī)學(xué)圖像分析的重要課題之一［1］。然而，通過(guò)重建神經(jīng)元結(jié)構(gòu)來(lái)獲得對(duì)人類(lèi)或動(dòng)物神經(jīng)系統(tǒng)功能連接性的深入了解是神經(jīng)解剖學(xué)者們面臨的重大挑戰(zhàn)［2］。電子顯微成像（Electron Microscope,EM）是神經(jīng)元結(jié)構(gòu)成像的重要設(shè)備，但其成像數(shù)量巨大，采用人工分割的方式雖然能獲得很好的分割效果，但這卻是一項(xiàng)枯燥、繁瑣且耗時(shí)耗力的任務(wù)［3］。Helmstaedter 等［4］使用20 000 h的人工僅從小鼠視網(wǎng)膜重建了大約1 000 個(gè)神經(jīng)元。自U-Net 在神經(jīng)元結(jié)構(gòu)分割中取得當(dāng)時(shí)最好的效果后，基于深度學(xué)習(xí)的分割算法成為醫(yī)學(xué)圖像分割算法的主流方法［5-6］。因此，設(shè)計(jì)一種基于深度學(xué)習(xí)的精確、快速、有效的神經(jīng)元結(jié)構(gòu)分割算法對(duì)于研究神經(jīng)系統(tǒng)功能有著重大意義。

自AlexNet 在2012年的ImageNet 競(jìng)賽 中取得冠軍后，卷積神經(jīng)網(wǎng)絡(luò)（Convolutional Neural Network,CNN）逐漸成為圖像特征提取的主流工具［7-8］。CNN的典型使用是在圖像分類(lèi)任務(wù)上，一些優(yōu)秀的模型如AlexNet［7］、VGG［9］、GoogleNet［10］、ResNet［11］等都會(huì)在網(wǎng)絡(luò)的最后加入一層全連接層，然后經(jīng)過(guò)softmax后獲得圖像的類(lèi)別信息。在生物醫(yī)學(xué)圖像分割任務(wù)中，期望的輸出應(yīng)當(dāng)是每個(gè)像素對(duì)應(yīng)的類(lèi)別。Long等［12］提出了全卷積網(wǎng) 絡(luò)（Fully Convolutional Network,FCN），主要思想是用卷積層代替全連接層，使得任意圖像大小的輸入成為可能，開(kāi)啟了一種端到端的訓(xùn)練模式。但是這種模型進(jìn)行語(yǔ)義分割時(shí)有一個(gè)主要的問(wèn)題：下采樣操作（比如pooling 層）會(huì)降低空間分辨率，削弱了空間“位置”信息。這個(gè)問(wèn)題在編碼器-解碼器的結(jié)構(gòu)（Encoder-Decoder）中得到有效解決。編碼器經(jīng)過(guò)下采樣逐漸減小空間維度提取高維信息，而解碼器逐漸恢復(fù)細(xì)節(jié)信息和空間維度?？衫胹hortcut conncetions 連接編碼解碼器不斷補(bǔ)全空間細(xì)節(jié)信息。U-Net 是生物醫(yī)學(xué)分割中的一種非常流行的編碼器-解碼器結(jié)構(gòu)［5］。在大規(guī)模數(shù)據(jù)集情況下，隨著網(wǎng)絡(luò)層數(shù)的加深，網(wǎng)絡(luò)的表達(dá)能力以及抽象能力也會(huì)隨之提高，并且已經(jīng)有實(shí)驗(yàn)證明，對(duì)于時(shí)間復(fù)雜度相同的兩種網(wǎng)絡(luò)結(jié)構(gòu)，深度較深的網(wǎng)絡(luò)性能會(huì)有所提升［13］。比如ResNet 在圖像識(shí)別中，ResNet101 的識(shí)別效果要優(yōu)于 ResNet34 和ResNet18［11］。同理，在語(yǔ)義分割任務(wù)中，一定情況下，模型越深分割效果越佳，DeepLab 利用ResNet101 作為骨架網(wǎng)絡(luò)，在PASCAL VOC 2012 以及MS-COCO數(shù)據(jù)集中取得了最好的結(jié)果［14-16］。然而，醫(yī)學(xué)圖像有一個(gè)比較嚴(yán)重的問(wèn)題：數(shù)據(jù)量小，深層模型在小樣本的情況下訓(xùn)練時(shí)容易出現(xiàn)過(guò)擬合的問(wèn)題。有研究者們發(fā)現(xiàn)，分組卷積能夠使過(guò)擬合問(wèn)題得到有效解決［17］。MobileNet 和Xception 成功地借鑒了這種思想，在圖像分類(lèi)和目標(biāo)檢測(cè)中能夠達(dá)到較好的實(shí)時(shí)性和性能［18-19］。

在醫(yī)學(xué)圖像分割中，損失函數(shù)的設(shè)計(jì)與選用尤為重要，常用的損失函數(shù)包括像素級(jí)交叉熵?fù)p失、Dice coefficient［20］、MSEloss 等。如何針對(duì)特定的樣本選用合適的損失函數(shù)，是目前研究的重點(diǎn)問(wèn)題之一。

基于此，本研究提出一種用于EM 圖像中神經(jīng)元結(jié)構(gòu)的分割深度學(xué)習(xí)模型：Group-Depth U-Net（G-D U-Net），實(shí)驗(yàn)結(jié)果表明該學(xué)習(xí)模型取得了比U-Net更好的分割效果。

1 方法與原理

1.1 G-D U-Net模型結(jié)構(gòu)

本研究的模型結(jié)構(gòu)見(jiàn)圖1，主要由4部分組成，即分組卷積、下采樣、特征圖融合和上采樣。

圖1 G-D U-Net模型示意圖Fig.1 Schematic representation of G-D U-Net model

1.1.1 分組卷積本研究的模型起始端分3 組卷積，因?yàn)檫^(guò)多的分組會(huì)增加卷積層的乘累加運(yùn)算（Multiply Accumulate,MAC）［17］，導(dǎo)致GPU 運(yùn)算量過(guò)大。第一組由輸入圖片直接進(jìn)行兩次下采樣，得到128 張feature map；第二組開(kāi)始由輸入圖片進(jìn)行卷積，得到32 張feature map，具體過(guò)程如圖1c 所示，然后進(jìn)行兩次下采樣，得到128 張feature map；第三組和第二組類(lèi)似，只是開(kāi)始時(shí)由輸入圖卷積得到64 張feature map。

MAC公式如下：

其中，h、w分別表示輸出特征圖的高度與寬度；c1、c2分別表示輸入的通道數(shù)和輸出的通道數(shù)；g表示分組卷積數(shù)。

1.1.2 特征圖融合模型起始端的3 組卷積得到了384 不同的feature map，將其進(jìn)行融合為下一步下采樣過(guò)程做準(zhǔn)備。

1.1.3 下采樣模型的下采樣過(guò)程是在原來(lái)的U-Net下采樣改進(jìn)的，原始的U-Net 一共下采樣4 次，由于前面進(jìn)行了分組卷積，為防止模型過(guò)深而過(guò)擬合，所以本研究增加了3 次下采樣，一共7 次下采樣。每次下采樣都是由3×3 的卷積核、ReLU 激活以及最大池化（Maxpooling）組成，具體過(guò)程見(jiàn)圖1a 所示。ReLU激活函數(shù)表示為：

1.1.4 上采樣為了恢復(fù)細(xì)節(jié)信息和空間維度，本研究的模型進(jìn)行了7 次上采樣，用的是雙線性插值（Bilinear Interpolation）［21］，實(shí)現(xiàn)端到端的訓(xùn)練，見(jiàn)圖1b 所示。同時(shí)利用shortcut conncetion 連接上采樣和下采樣，不斷補(bǔ)全空間細(xì)節(jié)信息，插值公式如下：

1.2 算法流程

提出的算法框圖如圖2所示，主要分兩大部分：（1）訓(xùn)練過(guò)程。為了訓(xùn)練本研究的分割模型，使用提供的標(biāo)準(zhǔn)EM 神經(jīng)元細(xì)胞圖像以及對(duì)應(yīng)的標(biāo)簽同時(shí)對(duì)訓(xùn)練集進(jìn)行適當(dāng)?shù)臄?shù)據(jù)增強(qiáng)，以提高模型的魯棒性。輸入訓(xùn)練的神經(jīng)元圖像，然后經(jīng)過(guò)G-D U-Net模型，得到一個(gè)預(yù)測(cè)分割圖像，與相對(duì)應(yīng)的標(biāo)簽做loss，然后經(jīng)過(guò)優(yōu)化器不斷優(yōu)化，直到loss 降到最低，模型訓(xùn)練完成，保存訓(xùn)練好的權(quán)重。（2）測(cè)試過(guò)程。輸入需要測(cè)試的EM 神經(jīng)元圖像，加載步驟1 訓(xùn)練好權(quán)重，然后經(jīng)過(guò)已經(jīng)訓(xùn)練好的G-D U-Net 模型，從而得到預(yù)測(cè)的分割圖像。

圖2 算法流程框圖Fig.2 Flowchart of the algorithm

2 結(jié)果

2.1 實(shí)驗(yàn)數(shù)據(jù)

本研究采用的數(shù)據(jù)集為ISBI 2012 提供的數(shù)據(jù)集［22-23］，訓(xùn)練數(shù)據(jù)是來(lái)自果蠅第一齡幼蟲(chóng)腹側(cè)神經(jīng)索（VNC）的連續(xù)部分透射電子顯微鏡（ssTEM）數(shù)據(jù)集的30 個(gè)部分，共30 張圖片。微管尺寸約為2.0 μm×2.0 μm×1.5 μm，分辨率為4×4×50 nm/像素。相應(yīng)的標(biāo)簽以二進(jìn)制的方式提供，即分割對(duì)象的像素為白色，其余像素為黑色（主要對(duì)應(yīng)于細(xì)胞膜）。為達(dá)到實(shí)驗(yàn)需求，本研究將其中的25 張圖像作為訓(xùn)練集，剩下的5張圖像作為測(cè)試集。為了避免過(guò)度擬合，本研究采用旋轉(zhuǎn)、翻轉(zhuǎn)和彈性變形等增強(qiáng)策略來(lái)擴(kuò)充數(shù)據(jù)集。經(jīng)過(guò)數(shù)據(jù)擴(kuò)充，本研究的訓(xùn)練數(shù)據(jù)達(dá)到800 張，其中90%作為訓(xùn)練集，10%作為驗(yàn)證集。同樣，將測(cè)試集擴(kuò)充為30 張。模型框架的實(shí)現(xiàn)是基于開(kāi)源軟件庫(kù)pytorch 0.4 版本，并在NVIDIA GeForce 1080Ti GPU 上采用端到端的方式訓(xùn)練。

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

本研究采用近年來(lái)比較常用的醫(yī)學(xué)圖像分割評(píng)價(jià)指標(biāo)，其中包含HD 系數(shù)、Dice 系數(shù)、Conformity 系數(shù)［24］、Sensitivity 系數(shù)和RECALL 系數(shù)。為了與UNet進(jìn)行對(duì)比，同樣使用了Warping Error［25］、Rand Error［26］、Pixel Error 這3 個(gè)評(píng)價(jià)指標(biāo)。每種指標(biāo)的介紹見(jiàn)表1。

表1 醫(yī)學(xué)圖像分割評(píng)價(jià)指標(biāo)介紹Tab.1 Introduction to the evaluation indexes of medical image segmentation

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

在ISBI 2012 數(shù)據(jù)集中的測(cè)試效果如圖3所示。從視覺(jué)感受上可以看出，本研究的模型在細(xì)節(jié)處理方面表現(xiàn)優(yōu)于原始的U-Net，對(duì)于神經(jīng)結(jié)構(gòu)內(nèi)部的一些不明確區(qū)域和復(fù)雜區(qū)域的分割效果更好。

圖3 分割結(jié)果的可視化Fig.3 Visualization of segmentation results

具體定性分析見(jiàn)表2～表4。表2為ISBI 2012 舊的評(píng)分指標(biāo)，從表2可知，G-D U-Net 模型在Rand Error 和Pixel Error 兩項(xiàng)指標(biāo)中優(yōu)于U-Net 模型，這是因?yàn)镚-D U-Net模型不僅加深了網(wǎng)絡(luò)層數(shù)，而且增加了分組卷積，從而補(bǔ)全了細(xì)節(jié)信息。表3為ISBI 2012 新的評(píng)分指標(biāo)，從表3可知，G-D U-Net 模型在新的評(píng)分指標(biāo)中也有優(yōu)異的表現(xiàn)。同樣，從表4可知，G-D U-Net 模型在5 種常用的醫(yī)學(xué)圖像分割指標(biāo)中表現(xiàn)比U-Net 模型好，從而驗(yàn)證了模型設(shè)計(jì)的有效性。

表2 不同模型分割效果的對(duì)比（a）Tab.2 Comparison of results of segmentation by different models(a)

表3 不同模型分割效果的對(duì)比（b）Tab.3 Comparison of results of segmentation by different models(b)

表4 不同模型分割效果的對(duì)比（c）Tab.4 Comparison of results of segmentation by different models(c)

2.3.1 模型的不同深度對(duì)分割性能的影響為了進(jìn)一步驗(yàn)證模型的深度對(duì)于分割效果的影響，本研究在G-D U-Net 模型原來(lái)7 次上下采樣的基礎(chǔ)上進(jìn)行改進(jìn)，得到兩個(gè)較淺的模型，一個(gè)是5 次上下采樣的模型：G-D U-Net1；另一個(gè)是6 次上下采樣的模型：G-D U-Net2。然后訓(xùn)練相同的次數(shù)對(duì)其進(jìn)行測(cè)試，得到的分割指標(biāo)如表5所示。從表5可知，模型的深度對(duì)分割效果有一定的影響，在一定范圍內(nèi)模型越深，分割效果會(huì)更好，但是網(wǎng)絡(luò)過(guò)深訓(xùn)練時(shí)容易過(guò)擬合，且容易出現(xiàn)計(jì)算機(jī)顯存不足的情況。

表5 不同深度的模型分割效果的對(duì)比Tab.5 Comparison of results of segmentation by different deep models

2.3.2 分組卷積對(duì)分割性能的影響同理，為了驗(yàn)證分組卷積對(duì)分割性能的影響，本研究在G-D U-Net的基礎(chǔ)上，移除第一組，得到G-D U-Net3；移除第一組和第二組卷積，保留第三組，得到新的模型G-D UNet4。測(cè)試效果如表6所示。模型起始端加入分割后，能夠在一定程度上提高分割效果。

表6 分組卷積對(duì)分割性能的影響Tab.6 Effect of group convolution on segmentation performance

2.3.3 不同的損失函數(shù)對(duì)分割性能的影響實(shí)驗(yàn)發(fā)現(xiàn)使用不同的損失函數(shù)對(duì)分割性能有一定的影響。本研究提出的G-D U-Net模型采用的是L1 loss，為了進(jìn)行對(duì)比，本研究把loss換成MSE loss和交叉熵?fù)p失函數(shù)（Cross Entropy loss,BCE loss），得到兩個(gè)新的模型，記為G-D U-Net5、G-D U-Net6。3 種loss 的分割效果圖如圖4所示，從視覺(jué)感受L1 loss最接近理想效果，其他兩種loss 分割后圖比較模糊。3 種損失函數(shù)的公式如下：

其中，lN表示第N個(gè)樣本中預(yù)測(cè)值與標(biāo)簽的誤差。

圖4 不同的損失函數(shù)對(duì)分割性能的影響Fig.4 Effect of different loss functions on segmentation performance

從以上的損失函數(shù)表達(dá)式可以看出，采用BCE loss 單獨(dú)評(píng)估每個(gè)像素矢量的分類(lèi)預(yù)測(cè)，然后對(duì)所有像素求平均值，基本上假定同等地對(duì)待圖像中每個(gè)像素點(diǎn)。如果各種類(lèi)在圖像分布不平衡時(shí)，這可能會(huì)是一個(gè)問(wèn)題，因?yàn)橛?xùn)練過(guò)程將受到最多的類(lèi)所支配。MSE loss 同樣也是利用預(yù)測(cè)的像素值與標(biāo)簽的像素值的loss 的平方求平均，所以也會(huì)出現(xiàn)上述問(wèn)題。因?yàn)榇朔指钊蝿?wù)是預(yù)測(cè)一張二值圖（細(xì)胞膜像素值為1，非膜像素為0），如果采用BCE loss 或者M(jìn)SE loss，預(yù)測(cè)結(jié)果就會(huì)出現(xiàn)如圖4c 和4d 上的“陰影”，而L1 loss 是將預(yù)測(cè)值與實(shí)際值誤差的絕對(duì)值作為loss，這樣預(yù)測(cè)值就能夠不斷接近真實(shí)值，不會(huì)出現(xiàn)“陰影”，因此L1 loss將更適用于此分割任務(wù)。

本研究對(duì)3 種不同損失函數(shù)的模型進(jìn)行測(cè)試，得到的分割指標(biāo)如表7所示。其中，G-D U-Net模型（采用L1 loss）在神經(jīng)元結(jié)構(gòu)的分割中表現(xiàn)最佳。

表7 不同的損失函數(shù)對(duì)分割性能的影響Tab.7 Effect of different loss functions on segmentation performance

3 總結(jié)與討論

本研究提出一種用于EM 圖像中神經(jīng)元結(jié)構(gòu)的分割深度學(xué)習(xí)模型G-D U-Net，在EM 神經(jīng)細(xì)胞的分割實(shí)驗(yàn)中，相比于原始的U-Net 模型，該模型獲得了較好的分割效果。同時(shí)本研究經(jīng)過(guò)大量的實(shí)驗(yàn)，驗(yàn)證了分組卷積、加深模型深度以及使用L1 loss 對(duì)分割效果的提升，為研究神經(jīng)結(jié)構(gòu)的生物醫(yī)學(xué)專(zhuān)家提供了一定的參考，也為其他生物圖像圖像分割領(lǐng)域提供一定的方法借鑒。