3D腦腫瘤分割的Dice損失函數(shù)的優(yōu)化

劉昊，王冠華，章強(qiáng)，李雨澤，陳慧軍

清華大學(xué) 醫(yī)學(xué)院 生物醫(yī)學(xué)工程系，北京 100084

引言

腦腫瘤是致命的疾病之一，全球每年發(fā)生原發(fā)性腦瘤的人數(shù)約為250000人，而腦瘤平均五年存活率為33%[1]。為了提高病人的生存率，盡早提供積極有效的治療有著重要的意義。僅對(duì)于臨床常見(jiàn)的腦膠質(zhì)瘤來(lái)說(shuō)，高級(jí)別膠質(zhì)瘤（High-Grade Glioma，HGG）病人的生存率較低。然而，通過(guò)對(duì)膠質(zhì)瘤的低級(jí)別（Low-Grade Glioma，LGG）早期階段進(jìn)行精確診斷，就可以較好的延長(zhǎng)病人的壽命[2-3]。因此盡早地對(duì)腦腫瘤進(jìn)行精確診斷十分重要。

磁共振成像（MRI）作為一種重要的臨床腦腫瘤成像技術(shù)可以為醫(yī)生的臨床診斷提供重要信息[4]。腦腫瘤在磁共振圖像上有著較為明顯的對(duì)比度，利于對(duì)腦腫瘤進(jìn)行準(zhǔn)確的分割。準(zhǔn)確地分辨腦腫瘤對(duì)于其診斷、治療和術(shù)后追蹤有著重要的意義，然而在臨床應(yīng)用中，目前主要使用手動(dòng)分割的方法來(lái)分割腦腫瘤，費(fèi)時(shí)費(fèi)力，且可重復(fù)性較差。因此全自動(dòng)高準(zhǔn)確度的腦腫瘤分割算法對(duì)臨床的準(zhǔn)確診斷和腦腫瘤的治療有著非常重要的意義。目前，國(guó)際上很多研究者進(jìn)行了磁共振腦腫瘤分割的研究。使用神經(jīng)網(wǎng)絡(luò)對(duì)醫(yī)學(xué)圖像進(jìn)行分割源自Ronneberger等[5]提出使用UNet對(duì)2D醫(yī)學(xué)圖像進(jìn)行分割，采用V-Net對(duì)3D圖像進(jìn)行分割并使用Dice系數(shù)作為損失函數(shù)，使得網(wǎng)絡(luò)可以優(yōu)化更加直接。Kamnitsas等[6]使用交叉熵?fù)p失函數(shù)作為損失函數(shù)，訓(xùn)練了deep medical，F(xiàn)CN[7]和 3D UNet的集成學(xué)習(xí)模型用來(lái)分割腦腫瘤[8]。Szegedy等[9]提出利用腫瘤區(qū)域之間的包含關(guān)系來(lái)進(jìn)行順序分割，并提出了類似于Inception Network的思路[10]，采用了3×3×1和1×1×3卷積核來(lái)替代3×3×3的卷積，既可以節(jié)約計(jì)算時(shí)間又可以減少GPU顯存的使用。Isensee等[10]提出采用通用3D UNet模型和額外的數(shù)據(jù)進(jìn)行分割，并使用了均值Sorensen Dice 損失函數(shù)+交叉熵?fù)p失函數(shù)作為損失函數(shù)。Huang等[11]使用DenseNet作為UNet編碼器并且利用3×3×3空洞卷積替代無(wú)空洞卷積以提高網(wǎng)絡(luò)視野域得到更好結(jié)果。Andriy Myronenko等人額外添加一個(gè)自動(dòng)解碼器對(duì)應(yīng)到原圖，然后使用原圖和真實(shí)圖（Ground Truth）所對(duì)應(yīng)的均值Jaccard損失函數(shù)+KL divergence損失函數(shù)作為損失函數(shù)以得到更好的結(jié)果[12]。

以上研究方法均優(yōu)化訓(xùn)練集標(biāo)簽所提供標(biāo)簽，本研究在端對(duì)端的卷積神經(jīng)網(wǎng)絡(luò)上提出了一種新的代價(jià)函數(shù)來(lái)提高分割的準(zhǔn)確率。該代價(jià)函數(shù)在融合均值Sorensen Dice代價(jià)函數(shù)和交叉熵函數(shù)的基礎(chǔ)上[10]，進(jìn)一步融合了全腫瘤區(qū)，核心腫瘤區(qū)這兩個(gè)重要臨床分割目標(biāo)的Dice系數(shù)。在實(shí)際測(cè)試中，本研究還對(duì)比了現(xiàn)有的其他代價(jià)函數(shù)，證實(shí)了新提出的代價(jià)函數(shù)有更高的分割準(zhǔn)確率。

1 方法

1.1 數(shù)據(jù)

本文采用Brats2018數(shù)據(jù)集[2-3]，包括285例病人，其中210例為高級(jí)別膠質(zhì)瘤（HGG），75例為低級(jí)別膠質(zhì)瘤（LGG）。每例病人的數(shù)據(jù)包括四個(gè)分辨率為240×240×155的剛性配準(zhǔn)過(guò)的3D磁共振序列（T1，T1Gd，T2和T2 Flair）圖像和相對(duì)應(yīng)的手動(dòng)標(biāo)記的“真實(shí)”分割結(jié)果（Ground Truth）。該數(shù)據(jù)集來(lái)自19個(gè)機(jī)構(gòu)使用不同的磁共振掃描儀所獲得的數(shù)據(jù)，標(biāo)記數(shù)據(jù)主要包括四個(gè)部分：背景（Background（BG），類型0）、壞死核心區(qū)和非增強(qiáng)核心區(qū)（Necrotic and non-Enhancing Tumor Core，NC&NETC，類型1）、水腫區(qū)（Edema，ED），類型2）和增強(qiáng)核心區(qū)（Enhanced Tumor，ET，類型 3），

1.2 數(shù)據(jù)預(yù)處理

數(shù)據(jù)的標(biāo)準(zhǔn)化：首先使用訓(xùn)練集的數(shù)據(jù)統(tǒng)計(jì)多對(duì)比度數(shù)據(jù)每個(gè)序列的均值和方差，然后利用統(tǒng)計(jì)的均值和方法分別對(duì)多對(duì)比度數(shù)據(jù)的每個(gè)圖像進(jìn)行零均值化和單位方差操作。

數(shù)據(jù)增值：首先采用隨機(jī)旋轉(zhuǎn)0°、隨機(jī)旋轉(zhuǎn)90°、180°、270°、隨機(jī)翻轉(zhuǎn)和0.9～1.1倍的隨機(jī)尺度變換。再采用隨機(jī)切割的方法分別將多對(duì)比度數(shù)據(jù)（配準(zhǔn)的大腦 3D T1、T1Gd、T2和 T2 Flair數(shù)據(jù)）分成多對(duì)比度的4×128×128×128的小塊數(shù)據(jù)[9]。

1.3 神經(jīng)網(wǎng)絡(luò)

本研究分割網(wǎng)絡(luò)使用的是基于改進(jìn)版本的端對(duì)端的3D UNet神經(jīng)網(wǎng)絡(luò)[13]作為分割網(wǎng)絡(luò)，主要結(jié)構(gòu)由編碼器和解碼器組成。其中編碼器主要由下采樣模塊（down sample）和帶殘差下采樣卷積模塊（dconv block）構(gòu)成，而解碼器主要由上采樣模塊（up sample）和帶殘差的上采樣卷積模塊（uconv block）組成。相比于傳統(tǒng)的3D UNet，不僅融合了前面的低尺度信息，同時(shí)使用上采樣模塊和卷積操作和解碼器上一層信息進(jìn)行疊加以融合更多的多尺度信息，分割層（seg layer）步幅為1的1×1×1的卷積操作，每層的輸出的通道數(shù)都為4。最后網(wǎng)絡(luò)通過(guò)softmax激活函數(shù)進(jìn)行輸出，整體網(wǎng)絡(luò)結(jié)構(gòu)如圖1所示，其中每個(gè)部分的詳細(xì)設(shè)計(jì)在1.3.1和1.3.2中詳細(xì)介紹。

圖1 腦腫瘤神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)分割圖

1.3.1 編碼器

編碼器主要由第一層卷積、下采樣模塊（down sample）和帶殘差的下采樣卷積模塊（dconv block）構(gòu)成。第一層卷積采用步幅（Stride）為1卷積核的個(gè)數(shù)（Number of Filter）為使用16的3×3×3的卷積，隨著每次下采樣卷積核個(gè)數(shù)進(jìn)行加倍。下采樣模塊為步幅為2的Max Pooling模塊。帶殘差的下采樣卷積模塊采用了類似ResNet的操作，主要由卷積，群標(biāo)準(zhǔn)化方法（Group Normalization，GN）[14]和激活函數(shù)構(gòu)成，具體如圖2a所示。

圖2 編碼器

1.3.2 解碼器

解碼器由上采樣模塊（Up Sample）和帶殘差下采樣卷積模塊（uconv block）構(gòu)成，其中上采樣模塊由三維的雙線性插值上采樣（3D Bilinear Upsampling）模塊組成。而uconv模塊如圖2b所示。為了更好地利用多尺度的信息，除了合并前面的低尺寸信息之外，Modified UNet額外增加了一條通道，其中分割層（seg layer）由步幅為1的1×1×1的卷積核構(gòu)成，輸出的卷積核個(gè)數(shù)為四個(gè)，如前所述，四個(gè)通道類別代表：背景1、壞死核心區(qū)2和非增強(qiáng)核心區(qū)3、水腫區(qū)和增強(qiáng)核心區(qū)4，其中通道1+2+3為全腫瘤區(qū)（Whole Tumor，WT）：通道1+通道2+通道3），通道 1+3 為核心腫瘤區(qū)（Tumor Core，TC）：通道 1+ 通道 3），見(jiàn)圖3。

圖3 通道說(shuō)明圖

1.4 損失函數(shù)

本文提出新的損失函數(shù)主要由三部分組成：交叉熵?fù)p失函數(shù)（Cross Entropy，CE））、均值 Sorensen Dice 損失函數(shù)[10]和直接優(yōu)化項(xiàng)的Sorensen Dice損失函數(shù)。交叉熵?fù)p失函數(shù)和Sorensen Dice[15]損失函數(shù)分別如公式(1)、(2)所示，其中p(x)為預(yù)測(cè)值而q(x)為相對(duì)應(yīng)的真值，其中p(x).sum()為p(x)中所有值相加，eps=1.0。

在本研究的實(shí)際問(wèn)題中，損失函數(shù)的第一部分是由均值Sorensen Dice損失函數(shù)Sorensen_Dice_NETC_ED_ET組成，具體見(jiàn)公式(3)所示。其中LNETC、LED和LET分別為其為NC&NETC、ED、ET的 Sorensen Dice損失函數(shù)構(gòu)成。

在實(shí)際問(wèn)題中，新提出的損失函數(shù)的第二部分主要由交叉熵?fù)p失函數(shù)構(gòu)成，具體見(jiàn)公式1。在融合了均值Sorensen Dice損失函數(shù)和交叉熵?fù)p失函數(shù)這兩部分之后，就可以得到損失函數(shù)Sorensen_Dice_NETC_ED_ET_CE，具體如公式(4)所示。

由于以上目標(biāo)函數(shù)在優(yōu)化的時(shí)，都未考慮TC、WT此兩項(xiàng)，而是間接優(yōu)化NETC、ED、ET，所以本文新提出的損失函數(shù)的第三部分主要由直接優(yōu)化項(xiàng)組成，試圖優(yōu)化臨床所需的TC和WT。即在公式(4)（Sorensen_Dice_NETC_ED_ET_CE）的基礎(chǔ)上加上以上分割目標(biāo)的Sorensen Dice損失函數(shù)得到最終的優(yōu)化函數(shù)Sorensen_Dice_NETC_ED_ET_WT_TC_CE，見(jiàn)公式(5)，其中LWT和LTC分別全腫瘤區(qū)，核心腫瘤區(qū)的Sorensen Dice損失函數(shù)，具體如圖3所示。

1.5 優(yōu)化器

本文使用Adam優(yōu)化器，權(quán)值衰減使用為1e-5，使用步衰減學(xué)習(xí)率（step learning rate）?？偣?450 個(gè)時(shí)期（epoch），初始學(xué)習(xí)率為1e-4，衰減率gamma等于0.5，前400個(gè)100個(gè)時(shí)期衰減一次。后面50個(gè)時(shí)期每25個(gè)epoch衰減一次。批尺寸（batch size）為 2。

1.6 訓(xùn)練與測(cè)試

本文使用 Pytorch 并在 NVIDIA Titan XP 12GB GPU 上完成訓(xùn)練與測(cè)試。我們將訓(xùn)練集HGG和LGG兩種情況分別按照4:1的比例分為訓(xùn)練集和測(cè)試集。本文將原尺寸圖片240×240×150填補(bǔ)（padding）到240×240×160輸入神經(jīng)網(wǎng)絡(luò)進(jìn)行測(cè)試。本研究的評(píng)價(jià)的指標(biāo)為：WT區(qū)為1:+2+3，TC區(qū)為1+3和ET區(qū)為3的Sorensen Dice系數(shù)。

作為對(duì)比，本文還比較了兩種現(xiàn)有的代價(jià)函數(shù)。首先對(duì)比了Sorensen Dice函數(shù)和Jaccard Dice函數(shù)所對(duì)應(yīng)的均值Sorensen Dice函數(shù)和均值 Jaccard Dice損失函數(shù)[13]（Jaccard_Dice_NETC_ED_ET：公式(7)。

其中LJaccard_Dice_NETC、LJaccard_Dice_WT和LJaccard_Dice_ET分別為NETC、WT、ET的Jaccard損失函數(shù)。

為了測(cè)試背景在均值Sorensen Dice損失函數(shù)中是否重要，還對(duì)比了Sorensen_Dice_NETC_ED_ET損失函數(shù)和添加背景的均值Sorensen Dice損失函數(shù)（Sorensen_Dice_NETC_ED_ET_BG），見(jiàn)公式 (8)。

2 結(jié)果

相比于交叉熵的損失函數(shù)和均值Sorensen Dice損失函數(shù)，本文提出的新?lián)p失函數(shù)Sorensen_Dice_NETC_ED_ET_WT_TC_CE在三項(xiàng)指標(biāo)（腫瘤區(qū)、腫瘤核心區(qū)和增強(qiáng)腫瘤區(qū)）上均在測(cè)試集上取得了最佳準(zhǔn)確率（表1），分別在腫瘤區(qū)、腫瘤核心區(qū)和增強(qiáng)腫瘤區(qū)這三個(gè)目標(biāo)區(qū)域的平均Dice系數(shù)分別達(dá)到：0.875、0.829、0.695。

表1 測(cè)試集結(jié)果

另外，在均值Jaccard Dice損失函數(shù)（Jaccard_Dice_NETC_ED_ET）和均值Sorensen Dice損失函數(shù)（Sorensen_Dice_NETC_ED_ET）損失函數(shù)的對(duì)比中，均值Sorensen Dice損失函數(shù)在WT、TC、ET三個(gè)指標(biāo)上都取得了領(lǐng)先的結(jié)果。盡管交叉熵?fù)p失函數(shù)（CE）在結(jié)果上不如均值Sorensen Dice損失函數(shù)（Sorensen_Dice_NETC_ED_ET），但是組合這兩者得到Sorensen_Dice_NETC_ED_ET_CE的結(jié)果相對(duì)于單獨(dú)使用兩者都可以得到一定的提升。而在背景項(xiàng)（BG）意義的測(cè)試中，加入背景項(xiàng)Sorensen_Dice_NETC_ED_ET_BG的分割準(zhǔn)確率低于不加背景項(xiàng)均值Sorensen Dice 損失函數(shù)。

本文提出的Sorensen_Dice_NETC_ED_ET_WT_TC_CE損失函數(shù)與Sorensen_Dice_NETC_ED_ET_CE和CE損失函數(shù)的收斂曲線對(duì)比圖，見(jiàn)圖4。可以看出，使用交叉熵?fù)p失函數(shù)（CE）的神經(jīng)網(wǎng)絡(luò)收斂曲線震蕩較大且收斂速度慢，而使用新提出的Sorensen_Dice_NETC_ED_ET_WT_TC_CE損失函數(shù)與Sorensen_Dice_NETC_ED_ET_CE代價(jià)函數(shù)的神經(jīng)網(wǎng)絡(luò)相比，在收斂速度基本相同的情況下獲得了更好的結(jié)果。

圖4 測(cè)試集收斂曲線

兩例病例的分割結(jié)果對(duì)比圖，見(jiàn)圖5，包括使用CE代價(jià)函數(shù)的分割圖（圖5c）、使用Sorensen_Dice_NETC_ED_ET_CE代價(jià)函數(shù)的分割圖（圖5d）、使用本文提出的新的Sorensen_Dice_NETC_ED_ET_WT_TC_CE代價(jià)函數(shù)的分割圖（圖5e），以及原圖（圖5a，選取Flair序列作為原圖）和人工分割的真實(shí)圖（圖5b）的對(duì)比。可以發(fā)現(xiàn)，使用三種不同的代價(jià)函數(shù)的神經(jīng)網(wǎng)絡(luò)所獲得的分割結(jié)果在整體WT上分割的結(jié)果較為相似。然而在TC和ET上的分割結(jié)果相差甚多，本文所提出的Sorensen_Dice_NETC_ED_ET_WT_TC_CE代價(jià)函數(shù)所獲得的分割圖不管是在ET和TC的分割上，還是在ED和NETC的分割上均最接近真實(shí)圖結(jié)果。

圖5 分割結(jié)果圖

3 討論

本文提出了融合均值Sorensen Dice函數(shù)、交叉熵函數(shù)和目標(biāo)分割區(qū)域代價(jià)函數(shù)的新?lián)p失函數(shù)，可直接在神經(jīng)網(wǎng)絡(luò)中優(yōu)化對(duì)于臨床非常重要的全腫瘤區(qū)、核心腫瘤區(qū)、增強(qiáng)腫瘤區(qū)評(píng)價(jià)指標(biāo)的Dice系數(shù)。相比于傳統(tǒng)的交叉熵?fù)p失函數(shù)和均值Sorensen/Jaccard損失函數(shù)，使用該損失函數(shù)的神經(jīng)網(wǎng)絡(luò)可以得到更加準(zhǔn)確的分割結(jié)果。本文中的方法與Rahman等[16]提出直接使用IOU損失函數(shù)作為損失函數(shù)來(lái)優(yōu)化神經(jīng)網(wǎng)絡(luò)的思想有相似之處，但本文采取不同的方式，將直接需要優(yōu)化項(xiàng)目標(biāo)項(xiàng)加入到已有的均值Sorensen損失函數(shù)。本研究的結(jié)果也表明，在分割不同種類的目標(biāo)時(shí)，針對(duì)主要分割目標(biāo)進(jìn)行代價(jià)函數(shù)的設(shè)計(jì)有可能得到更準(zhǔn)確地分割結(jié)果。

另外，本研究結(jié)果表明，在均值Sorensen Dice函數(shù)中加入背景優(yōu)化，反而會(huì)導(dǎo)致目標(biāo)分割的準(zhǔn)確性下降。可能的原因是，在Brats2018的數(shù)據(jù)集中，壞死核心區(qū)和非增強(qiáng)核心區(qū)，水腫區(qū)，增強(qiáng)核心區(qū)樣本之間在訓(xùn)練集的像素比例為：25%、57%和18%，相對(duì)比較平衡，而B(niǎo)G和NETC+ED+ET像素比例接近90，導(dǎo)致神經(jīng)網(wǎng)絡(luò)更加傾向于背景信息的優(yōu)化，導(dǎo)致其他項(xiàng)在損失函數(shù)中變得不重要，最終嚴(yán)重影響了最終分類的準(zhǔn)確率。

在WT、TC、ET這三個(gè)分割目標(biāo)中，可以發(fā)現(xiàn)不管使用哪種代價(jià)函數(shù)，ET的分割準(zhǔn)確率都是最低的。這不僅是因?yàn)镋T相對(duì)于其他組織的對(duì)比度較差，而且由于圖像采自于不同的設(shè)備，導(dǎo)致不同圖像的ET對(duì)比度不盡相同，需要網(wǎng)絡(luò)有很強(qiáng)的泛化能力才能得到精確的結(jié)果，我們接下來(lái)的工作將會(huì)集中在如何提高網(wǎng)絡(luò)的泛化能力，使得網(wǎng)絡(luò)能夠更好地契合不同機(jī)器采集的圖像以得到更好的結(jié)果。