于智洪,李菲菲

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

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

解決“域偏移”問題的直接方法是通過獲取大量多樣的含標簽數(shù)據(jù)集進行模型訓(xùn)練,但該方式耗時耗力并且需要大量資金。其次,域適應(yīng)(Domain Adaptation,DA)[5]方法也被廣泛用于解決“域偏移”問題。該方法主要通過訓(xùn)練源域(Source Domain)模型,并使用目標域(Target Domain)中的部分數(shù)據(jù)進行模型再訓(xùn)練,通過該方式將目標域中所涉及的特征信息用于目標域評估。但在醫(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)表征學(xué)習(xí)(Representation Learning,RL)[8],即學(xué)習(xí)域不變(Domain-invariant)表示,或者是將特征分解為特定和公共兩部分表達式;3)基于梯度的元學(xué)習(xí)方法[9]。本文采用基于梯度的元學(xué)習(xí)方法來獲得域不變表示以達到更好的域泛化能力。該方法將源域分為元訓(xùn)練集和元測試集,通過這兩個集合不斷模擬域偏移過程學(xué)習(xí)提取域不變特征。先前方法大部分是在監(jiān)督學(xué)習(xí)下完成,例如文獻[10]利用了全局類對齊和局部樣本聚類,文獻[11]使用形狀感知約束,但醫(yī)學(xué)圖像的像素級標注費時費力,所以在數(shù)據(jù)稀少的環(huán)境中可能無法獲取準確的域不變特征。在上述基于元學(xué)習(xí)的分割方法[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)搜索,使其可以與元學(xué)習(xí)更好地配合,以探索出準確的域不變特征,實現(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ù)。而元學(xué)習(xí)[15]致力于從經(jīng)驗中進行學(xué)習(xí),該方法可以從少量數(shù)據(jù)進行學(xué)習(xí),并應(yīng)用到新的任務(wù)之后。將二者結(jié)合可產(chǎn)生一個利用少量訓(xù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è)計了一個新訓(xùn)練策略,首先針對不同域的稀少數(shù)據(jù)進行架構(gòu)搜索以找到較合適的分割網(wǎng)絡(luò),使分割結(jié)果更準確。其次利用基于梯度的半監(jiān)督元學(xué)習(xí)方法,將源域分為元訓(xù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ù)作為源域訓(xùn)練集,選擇一個新域作為目標域數(shù)據(jù)集以評估泛化能力。如圖1所示,本文首先利用m%的標簽數(shù)據(jù)和神經(jīng)架構(gòu)搜索進行分割網(wǎng)絡(luò)架構(gòu)的搜索與確定,在分割網(wǎng)絡(luò)確定后使用半監(jiān)督元學(xué)習(xí)訓(xùn)練[13]策略進一步探索域不變表示。首先,源域數(shù)據(jù)被分為元訓(xùn)練集tr和元測試集te,在每次迭代過程中,元訓(xùn)練集tr通過Dice損失函數(shù)meta-train對β和φ進行更新

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

(1)

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

(2)

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

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

由于在醫(yī)學(xué)圖像分割中U型網(wǎng)絡(luò)架構(gòu)是較主流的方法,因此本文基于NAS-UNet[17]和UNet++[18]將搜索空間確定為U型結(jié)構(gòu),由上采樣、下采樣以及正常采樣3種單元構(gòu)成,從不同尺度中學(xué)習(xí)最優(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 元學(xué)習(xí)訓(xùn)練策略

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 元訓(xùn)練與元測試目標函數(shù)

首先,元訓(xùn)練目標函數(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)化器,學(xué)習(xí)率為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é)果,其訓(xùn)練數(shù)據(jù)集包含2%或者5%的標簽數(shù)據(jù)和全部無標簽數(shù)據(jù),而基準模型僅使用2%或者5%的標簽數(shù)據(jù)進行訓(xùn)練。粗體數(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進行自動分割,使用元學(xué)習(xí)進行訓(xùn)練取得了更好的泛化效果,精度提升至少3%。針對使用了圖像增強的半監(jiān)督學(xué)習(xí)方法SDNet[26],采用rank損失的LDDG[23]以及基于梯度下降的元學(xué)習(xí)分割方法SAML[11],本文方法依然擁有2%以上的精度優(yōu)勢。以上實驗進一步證明本文方法將元學(xué)習(xí)與神經(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)搜索和元學(xué)習(xí)來進行泛化醫(yī)療分割方法研究。首先利用U型搜索空間,針對不同源域進行上、下以及正常單元結(jié)構(gòu)搜索,使其分割網(wǎng)絡(luò)更適合本源域的分類,繼而利用元學(xué)習(xí)進行泛化學(xué)習(xí),使其更好地適應(yīng)未知目標域數(shù)據(jù)以達到良好的泛化效果。目前,本文方法在M&Ms數(shù)據(jù)集上擁有良好性能,但該網(wǎng)絡(luò)依然存在缺陷,例如當源域變化時,本文方法需要重新搜索網(wǎng)絡(luò)結(jié)構(gòu),耗費較多時間。后續(xù)工作將研究利用元學(xué)習(xí)使分割結(jié)構(gòu)也納入具有快速泛化能力的新方法,從而節(jié)約時間和成本。