基于元學習和神經(jīng)架構(gòu)搜索的半監(jiān)督醫(yī)學圖像分割方法

于智洪,李菲菲

(上海理工大學 光電信息與計算機工程學院,上海 200093)

隨著深度學習的發(fā)展[1],盡管醫(yī)學圖像分割方法[2]已有較大進展,但當將訓練好的模型用于未知數(shù)據(jù)集時可能引起性能驟降,這主要是由病人種群、掃描儀器和掃描采集設(shè)置不同而引起,該問題被稱為“域偏移(Domain Shift,DS)”[3]。其中,性別、年齡、種族特征等因素導致潛在病理和解剖結(jié)構(gòu)有所差異[4]。掃描儀器和掃描采集設(shè)置會影響醫(yī)學圖像的亮度、對比度等特征,這也給醫(yī)學圖像任務(wù)造成困擾。

解決“域偏移”問題的直接方法是通過獲取大量多樣的含標簽數(shù)據(jù)集進行模型訓練,但該方式耗時耗力并且需要大量資金。其次,域適應(yīng)(Domain Adaptation,DA)[5]方法也被廣泛用于解決“域偏移”問題。該方法主要通過訓練源域(Source Domain)模型,并使用目標域(Target Domain)中的部分數(shù)據(jù)進行模型再訓練,通過該方式將目標域中所涉及的特征信息用于目標域評估。但在醫(yī)療診斷現(xiàn)實應(yīng)用中,任何未知的目標域數(shù)據(jù)集均可能遇到“域偏移”問題,因此域適應(yīng)方法仍然無法直接應(yīng)用于現(xiàn)實社會生活?？紤]到醫(yī)療診斷在現(xiàn)實社會中的應(yīng)用問題,一個更加具有挑戰(zhàn)性的方法被提出——域泛化(Domain Generalization,DG)[6],該方法不使用目標域中的任何先驗知識,而是通過對源域數(shù)據(jù)集的探索來更好地泛化到未知醫(yī)療任務(wù)中。因此,域泛化更具有現(xiàn)實應(yīng)用意義和挑戰(zhàn)性。

就域泛化問題而言,目前可實現(xiàn)該目標的方法分為:1)源域數(shù)據(jù)增強[7];2)表征學習(Representation Learning,RL)[8],即學習域不變(Domain-invariant)表示,或者是將特征分解為特定和公共兩部分表達式;3)基于梯度的元學習方法[9]。本文采用基于梯度的元學習方法來獲得域不變表示以達到更好的域泛化能力。該方法將源域分為元訓練集和元測試集,通過這兩個集合不斷模擬域偏移過程學習提取域不變特征。先前方法大部分是在監(jiān)督學習下完成,例如文獻[10]利用了全局類對齊和局部樣本聚類,文獻[11]使用形狀感知約束,但醫(yī)學圖像的像素級標注費時費力,所以在數(shù)據(jù)稀少的環(huán)境中可能無法獲取準確的域不變特征。在上述基于元學習的分割方法[12-13]中,多數(shù)采用固定的U型結(jié)構(gòu)網(wǎng)絡(luò)進行分割,卻并未考慮到針對不同源域任務(wù)利用更適合的分割結(jié)構(gòu)可以獲得更好的結(jié)果。

因此,本文采用了神經(jīng)架構(gòu)搜索(Neural Architecture Search,NAS)[14]對特定源域任務(wù)進行分割結(jié)構(gòu)搜索,使其可以與元學習更好地配合,以探索出準確的域不變特征,實現(xiàn)良好的泛化能力。神經(jīng)架構(gòu)搜索由搜索空間、搜索策略和評估技巧3部分組成,由于設(shè)計一個性能良好的網(wǎng)絡(luò)結(jié)構(gòu)需要大量人力,神經(jīng)架構(gòu)搜索可以通過超參數(shù)的優(yōu)化更新來自動獲取較合適的針對特定任務(wù)的架構(gòu),但該過程需要大量標簽數(shù)據(jù)。而元學習[15]致力于從經(jīng)驗中進行學習,該方法可以從少量數(shù)據(jù)進行學習,并應(yīng)用到新的任務(wù)之后。將二者結(jié)合可產(chǎn)生一個利用少量訓練數(shù)據(jù)和迭代次數(shù)并且可以自適應(yīng)網(wǎng)絡(luò)結(jié)構(gòu)的方法。

本文在DGNet[13]基礎(chǔ)上提出了更加靈活準確的域泛化網(wǎng)絡(luò)?；谏窠?jīng)架構(gòu)搜索的自適應(yīng)性[16],本文設(shè)計了一個新訓練策略,首先針對不同域的稀少數(shù)據(jù)進行架構(gòu)搜索以找到較合適的分割網(wǎng)絡(luò),使分割結(jié)果更準確。其次利用基于梯度的半監(jiān)督元學習方法,將源域分為元訓練集和元測試集,利用上述分割結(jié)果不斷探索合適的域不變表示。實驗結(jié)果表明,與DGNet以及一些主流網(wǎng)絡(luò)相比,本文方法有效提高了域泛化能力,擁有良好的性能。

1 本文方法

1.1 算法流程

源域數(shù)據(jù)由標簽數(shù)據(jù)和無標簽數(shù)據(jù)兩部分組成。本文選擇所有無標簽數(shù)據(jù)和m%(m=0,1,2,…,100)的標簽數(shù)據(jù)作為源域訓練集,選擇一個新域作為目標域數(shù)據(jù)集以評估泛化能力。如圖1所示,本文首先利用m%的標簽數(shù)據(jù)和神經(jīng)架構(gòu)搜索進行分割網(wǎng)絡(luò)架構(gòu)的搜索與確定,在分割網(wǎng)絡(luò)確定后使用半監(jiān)督元學習訓練[13]策略進一步探索域不變表示。首先,源域數(shù)據(jù)被分為元訓練集tr和元測試集te,在每次迭代過程中,元訓練集tr通過Dice損失函數(shù)meta-train對β和φ進行更新

圖1 本文算法整體框架Figure 1. Overall framework of the proposed algorithm

(1)

其中,λ代表元訓練tr更新過程中的學習率。元測試te使用更新過后的(β′,φ′)進行再訓練,使元訓練tr和元測試te過程中的損失和最小[16]

(2)

無論是源域數(shù)據(jù)集還是未知數(shù)據(jù)集,通過該方法都可以擁有良好的泛化性能。最后,使用訓練得到的最佳模型在新域數(shù)據(jù)集上進行評估。

1.2 分割網(wǎng)絡(luò)架構(gòu)搜索

由于在醫(yī)學圖像分割中U型網(wǎng)絡(luò)架構(gòu)是較主流的方法,因此本文基于NAS-UNet[17]和UNet++[18]將搜索空間確定為U型結(jié)構(gòu),由上采樣、下采樣以及正常采樣3種單元構(gòu)成,從不同尺度中學習最優(yōu)特征表示,獲得較合適的子網(wǎng)絡(luò)。

1.2.1 搜索空間

如圖2所示,該搜索空間深度為5,不同尺度的特征映射通過跳躍連接進行信息融合。其內(nèi)部搜索空間結(jié)構(gòu)由先前的兩個特征圖作為輸入,進行上采樣、下采樣、正常采樣操作。具體下采樣單元結(jié)構(gòu)如圖2左下所示,下采樣與正常采樣單元結(jié)構(gòu)相同。圖2可以用有向無環(huán)圖(Directed Acyclic Graph,DAG)來解釋,xk-2和xk-1為兩個輸入,0、1、2、3為中間4個潛在的節(jié)點,xk為輸出。信息通過特征映射(i,j)(i

表1 候選操作總結(jié)

圖2 分割網(wǎng)絡(luò)的搜索空間Figure 2. The search space of segmentation network

與DARTS[19]相似,每條邊都由超參數(shù)α(i,j)來權(quán)衡重要性,本文依據(jù)O的維度,針對3種單元設(shè)置了3個不同的權(quán)重矩陣,使其從離散空間變?yōu)檫B續(xù)空間,并通過對應(yīng)的損失函數(shù)和梯度下降進行參數(shù)更新,最后使用Softmax函數(shù)對具體邊應(yīng)選取的操作進行計算

(3)

1.2.2 搜索策略

1.3 元學習訓練策略

1.3.1 特征提取網(wǎng)絡(luò)

圖3 特征網(wǎng)絡(luò)結(jié)構(gòu)Figure 3. The structure of feature network

(4)

其中,c是域標簽,超參數(shù)λrank=0.1,λKL=0.1,λrec=1.0,λcls=1.0。

1.3.2 元訓練與元測試目標函數(shù)

首先,元訓練目標函數(shù)由兩部分組成

(5)

當標簽數(shù)據(jù)可獲得時,λDice=5。在元測試階段,本文致力于獲得更加準確的分割掩碼以及保持s、d和Z之間的糾纏關(guān)系。因此,元測試階段的目標函數(shù)如式(6)所示。

(6)

2 實驗

2.1 實驗數(shù)據(jù)集

本文實驗使用多中心、多供應(yīng)商和多疾病心臟圖像分割(M&Ms)數(shù)據(jù)集[4]來驗證泛化能力。該數(shù)據(jù)集包含320名受試者,他們在3個不同國家的6個臨床中心接受掃描,并且使用了西門子、飛利浦、GE和佳能的核磁共振儀器,即A、B、C、D這4個不同域,其中域A包含95個受試者,域B包含125個受試者,域C和域D各包含50個受試者。

2.2 實驗配置

實驗使用Pytorch[24]框架和一塊NVidia 2080 Ti GPU,所有模型均使用Adam優(yōu)化器,學習率為2×10-5,批尺寸為2,一共迭代150次,重采樣率為1.2。在半監(jiān)督設(shè)置中,本文使用特定百分比的受試者作為標簽數(shù)據(jù)和所有無標簽數(shù)據(jù)。本文使用Dice和Hausdorff Distance作為評價指標。

2.3 實驗結(jié)果與討論

表2～表5展示了Dice和Hausdorff Distance以及它們的標準差結(jié)果,其訓練數(shù)據(jù)集包含2%或者5%的標簽數(shù)據(jù)和全部無標簽數(shù)據(jù),而基準模型僅使用2%或者5%的標簽數(shù)據(jù)進行訓練。粗體數(shù)字表示最佳性能。由表得知本文方法可以獲得較好的泛化性能。就表2和表3而言,在Dice指標下,無論是5%還是2%的標簽數(shù)據(jù),本文方法有超過一半的精度均好于DGNet,并且高出約2%。雖然5%標簽的精度平均值與DGNet持平,但是在2%標簽數(shù)據(jù)下,相比DGNet以及其他主流方法,本文方法至少超出1.2%。就Hausdorff Distance指標而言,當其距離越小,表明效果越好。從表4～表5可以看出,本文方法無論目標域為何,與所有主流方法相比,均可以取得最佳效果,超出至少2%。

表2 M&Ms在含5%標簽數(shù)據(jù)Dice上的結(jié)果

表3 M&Ms在含2%標簽數(shù)據(jù)Dice上的結(jié)果

表4 M&Ms在含5%標簽數(shù)據(jù)Hausdorff distance上的結(jié)果

表5 M&Ms在含2%標簽數(shù)據(jù)Hausdorff distance上的結(jié)果

與自動分割網(wǎng)絡(luò)nnUNet[25]相比,本方法使用NAS進行自動分割,使用元學習進行訓練取得了更好的泛化效果,精度提升至少3%。針對使用了圖像增強的半監(jiān)督學習方法SDNet[26],采用rank損失的LDDG[23]以及基于梯度下降的元學習分割方法SAML[11],本文方法依然擁有2%以上的精度優(yōu)勢。以上實驗進一步證明本文方法將元學習與神經(jīng)架構(gòu)搜索結(jié)合的有效性和良好的泛化能力。

當域A作為目標域時,圖4為在5%和2%標簽條件下的分割結(jié)果以及其真實值,可以看出本文方法具有良好的泛化能力。表6為其在5%標簽下搜索出的最優(yōu)單元結(jié)構(gòu)。由圖2可知,一共有4個節(jié)點,每個節(jié)點由前兩個輸出與操作卷積進行連接,然后再進行下一步輸出。通過該結(jié)構(gòu)可使網(wǎng)絡(luò)擁有較好的分割結(jié)果,從而產(chǎn)生較好的泛化能力。

表6 M&Ms在含5%標簽數(shù)據(jù)上的單元結(jié)構(gòu)

圖4 分割可視化結(jié)果Figure 4. The visualization results of segmentation

3 結(jié)束語

本文利用基于梯度下降的神經(jīng)架構(gòu)搜索和元學習來進行泛化醫(yī)療分割方法研究。首先利用U型搜索空間,針對不同源域進行上、下以及正常單元結(jié)構(gòu)搜索,使其分割網(wǎng)絡(luò)更適合本源域的分類,繼而利用元學習進行泛化學習,使其更好地適應(yīng)未知目標域數(shù)據(jù)以達到良好的泛化效果。目前,本文方法在M&Ms數(shù)據(jù)集上擁有良好性能,但該網(wǎng)絡(luò)依然存在缺陷,例如當源域變化時,本文方法需要重新搜索網(wǎng)絡(luò)結(jié)構(gòu),耗費較多時間。后續(xù)工作將研究利用元學習使分割結(jié)構(gòu)也納入具有快速泛化能力的新方法,從而節(jié)約時間和成本。