一種Unet圖像分割模型的聯(lián)邦蒸餾優(yōu)化算法

林子謙,樊重俊,王 琪

(上海理工大學(xué) 管理學(xué)院,上海 200093)

1 引 言

深度學(xué)習(xí)已成功地應(yīng)用于醫(yī)學(xué)圖像的病理檢測和診斷中[1],它可以在沒有醫(yī)學(xué)專家輸入的標(biāo)準(zhǔn)圖像處理規(guī)則的情況下原始圖像中學(xué)習(xí)并提供高精度的輸出.但這需要大量不同的數(shù)據(jù)才能發(fā)揮作用,若數(shù)據(jù)集差異很小或者疾病數(shù)據(jù)集很少時,神經(jīng)網(wǎng)絡(luò)是沒辦法有良好的泛化能力的.此外,在訓(xùn)練醫(yī)學(xué)影像模型中也可能依據(jù)機(jī)構(gòu)偏見相關(guān)的混雜因素,而不是將模型建立在所評估的病理基礎(chǔ)上,該模型在檢驗同一機(jī)構(gòu)的保留數(shù)據(jù)時,雖然準(zhǔn)確度較高,但對外部機(jī)構(gòu)的了解較少,不能很好地擴(kuò)展到外部機(jī)構(gòu).理論上,最有效的方法是將有關(guān)患者的所有數(shù)據(jù)集中起來共享到一個中心,并進(jìn)行訓(xùn)練,能夠產(chǎn)生良好的性能.但此方法存在許多挑戰(zhàn),首先,如果患者的數(shù)據(jù)占用了大量的存儲,共享這些數(shù)據(jù)將是一種麻煩;其次,共享患者數(shù)據(jù)常常有法律或道德上的障礙,使研究者無法獲得所有數(shù)據(jù);第三,患者數(shù)據(jù)很有價值,機(jī)構(gòu)可能不愿意分享數(shù)據(jù).

因此聯(lián)邦學(xué)習(xí)[2-4]孕育誕生,聯(lián)邦學(xué)習(xí)主要是借助客戶端-服務(wù)器的基本架構(gòu),實現(xiàn)多客戶端在服務(wù)器的協(xié)調(diào)下,

對訓(xùn)練的模型參數(shù)進(jìn)行聚合與傳遞,并訓(xùn)練出一個高性能的共享模型.這一概念的正式提出是由谷歌[5]HB Mcmahan等人,他們通過使用邊緣服務(wù)器架構(gòu),將聯(lián)邦學(xué)習(xí)用于智能手機(jī)上的語言預(yù)測模型的更新.隨后,發(fā)現(xiàn)聯(lián)邦學(xué)習(xí)有望打破在醫(yī)療領(lǐng)域中的隱私問題,許多學(xué)者也開始通過聯(lián)邦學(xué)習(xí)算法設(shè)計了在腦腫瘤分離[6,7]、乳房密度分類[8]和功能性磁共振成像分析[9]等醫(yī)學(xué)相關(guān)研究.

盡管聯(lián)邦學(xué)習(xí)在醫(yī)療領(lǐng)域有一定的相關(guān)研究,但是在實際應(yīng)用中經(jīng)常會面臨以下問題:1)由于經(jīng)典醫(yī)學(xué)圖像分割模型結(jié)構(gòu)復(fù)雜會產(chǎn)生大量的模型參數(shù),若發(fā)送如此大量的模型參數(shù),服務(wù)器將會產(chǎn)生巨大的通信開銷,并且這樣的通信開銷會隨著參與方數(shù)量和迭代輪次的增加而增加.此外,如果聯(lián)邦學(xué)校使用同態(tài)加密作為安全方案,那么加密后的密文數(shù)據(jù)在傳輸效率過程中會產(chǎn)生更多的通信開銷.2)現(xiàn)有醫(yī)學(xué)圖像分割的訓(xùn)練數(shù)據(jù)一般都是遵循獨立同分布的(Identically Independently Distributions,IID),但是應(yīng)用到聯(lián)邦學(xué)習(xí)當(dāng)中時大多數(shù)都是非獨立同分布(non-IID)的,甚至有些醫(yī)療機(jī)構(gòu)會傳遞惡意的模型,這必然會造成全局模型的梯度爆炸與準(zhǔn)確度下降的問題.

解決通信開銷方面,一般分為通信中的總時耗與通信數(shù)據(jù)總量這兩個問題,主要的解決方法有降低模型的更新頻率、模型的壓縮、客戶端的選擇.在降低模型的更新頻率方法中,Zhou等[9]提出重疊聯(lián)邦平均算法,通過集合分層計算策略、數(shù)據(jù)補(bǔ)償機(jī)制與NAG的算法,實現(xiàn)通信成本的降低;在模型的壓縮方法中,Yang等[10]以非線性的方式劃分空間,提出了一種基于cosine函數(shù)的非線性量化方案,不需要誤差反饋等額外梯度恢復(fù)信息來調(diào)整梯度;在客戶端的選擇方法中,Nishio T[11]提出了一種聯(lián)邦客戶選擇算法,該算法根據(jù)累計有效參與值,選擇模型迭代效率最高的客戶端進(jìn)行聚合更新,以此優(yōu)化整個聯(lián)邦學(xué)習(xí)算法的收斂效率,進(jìn)而降低通信代價.

解決數(shù)據(jù)的非獨立同分布主要有3種情況:特征分布偏差、標(biāo)簽分布偏差、數(shù)據(jù)量的偏差.而在圖像分割中經(jīng)常產(chǎn)生的偏差莫過于不同機(jī)構(gòu)的標(biāo)記不同,導(dǎo)致的特征分布的偏差,與不同數(shù)量的患者數(shù)據(jù)導(dǎo)致的偏差.因此針對非獨立同分布數(shù)據(jù)的聯(lián)邦學(xué)習(xí)問題也引起了研究者的關(guān)注,其中一個方向是優(yōu)化模型聚合過程,Wang等人[12]提出了一種歸一化的平均方法辦法,該辦法消除了目標(biāo)的不一致性,同時保留了異質(zhì)性的快收斂.另一個方向是通過知識蒸餾的方法,FedMA[13]完善了通過匹配和平均具有類似特征的隱藏元素來完善聚合過程具有類似特征簽名的隱藏元素,將知識提煉整合到聯(lián)邦學(xué)習(xí)中的想法也已經(jīng)得到了驗證[14].

本文設(shè)計一種針對圖像分割領(lǐng)域的聯(lián)邦蒸餾優(yōu)化算法(Unet-Federated Distillation Optimize,U-FedDO)算法.具體而言,首先在聯(lián)邦平均算法(Federated Averaging,FedAvg)的基礎(chǔ)上使用教師-學(xué)生模型,其中在參與方中有一個經(jīng)典的Unet模型作為教師模型,還有一個輕量的U-net模型結(jié)構(gòu)作為學(xué)生模型,通過教師-學(xué)生的框架保證本地參與方性能的穩(wěn)定;然后本文在服務(wù)器上還部署調(diào)節(jié)器,并設(shè)定設(shè)置相關(guān)的約束條件與權(quán)重設(shè)置機(jī)制.此外,本文還應(yīng)用在真實數(shù)據(jù)上設(shè)計了IID數(shù)據(jù)與non-IID數(shù)據(jù)下的實驗,結(jié)果充分地表明了該模型在醫(yī)學(xué)圖像分割中的有效性與可行性.

本文的組織結(jié)構(gòu)如下,第2節(jié)討論了相關(guān)工作,第3節(jié)描述了U-FedDO算法的具體細(xì)節(jié),第4節(jié)給出了實驗結(jié)果并討論,在第5節(jié)對本文進(jìn)行總結(jié)與討論.

2 相關(guān)工作

U-FedDO是針對圖像分割所設(shè)計的聯(lián)邦集成模型,其中包括了醫(yī)學(xué)圖像分割模型,聯(lián)邦優(yōu)化算法模型,聯(lián)邦蒸餾算法.下面對3個模型方法簡單介紹:

2.1 醫(yī)學(xué)圖像分割模型

醫(yī)學(xué)圖像分割是一種能夠?qū)θ梭w某個組織器官的邊界、形狀進(jìn)行測量和標(biāo)注的技術(shù).傳統(tǒng)的圖像分割方法通常是基于區(qū)域分割法和基于邊界分割方法.2015年,全卷積神經(jīng)網(wǎng)絡(luò)(Fully Convolutional Networks,FCNs)以一種端到端的形式進(jìn)行語義分割[15].隨后出的U-Net具有靈活、簡單并且在樣本數(shù)據(jù)集較少的情況下能獲得很好的分割效果,目前已成為醫(yī)療圖像分割的一類主流方法[16].

U-net結(jié)構(gòu)主要由兩部分組成:下采樣與上采樣.下采樣是當(dāng)圖像輸入Unet模型后采用一定的卷積用來圖像深度的提升與圖像特征的提取,接著使用池化操作用于減小圖像的面積.在對圖像進(jìn)行4層卷積池化后,進(jìn)行圖像全連接.上采樣是因為圖像的分割結(jié)果需要與輸入圖像相匹配,需要對最小圖像面積進(jìn)行擴(kuò)充,將圖像通過反卷積操作,并將上方相同層次的圖像合并,保留圖像原始特征.最后,使用一定的卷積核進(jìn)行圖像卷積,提取圖像特征.此外,為了防止傳統(tǒng)神經(jīng)網(wǎng)絡(luò)在信息傳遞過程中出現(xiàn)的梯度消失,在Unet結(jié)構(gòu)里的編碼器與解碼器之間加入了跳躍連接操作.后來在醫(yī)學(xué)圖像分割領(lǐng)域中,出現(xiàn)許多改進(jìn)Unet提升性能與應(yīng)用場景的圖像分割模型,主要包括編、解碼器結(jié)果的改進(jìn)、損失函數(shù)的改進(jìn)、瓶頸模塊的改進(jìn)、增強(qiáng)數(shù)據(jù)流路徑、自動搜索結(jié)果等.例如在解碼器中引入注意力機(jī)制的Attention Unet模型,抑制一些不相關(guān)分割區(qū)域所產(chǎn)生的影響,從而提升性能,或者是net++模型在跳躍連接的方式中加入不同深淺特征的融合,還有像Unet-GAN模型是從內(nèi)部結(jié)構(gòu)中提出的適應(yīng)不同提供商圖像的適配器.然而,目前基于Unet提升性能的方法主要集中于對Unet結(jié)構(gòu)的本身,卻很少有學(xué)者通過解決醫(yī)學(xué)圖像數(shù)據(jù)的孤島的問題,從而提升分割模型的性能.

2.2 聯(lián)邦優(yōu)化算法

U-FedDO算法主要是基于在FedAvg算法上進(jìn)行改進(jìn).傳統(tǒng)的FedAvg是指參與方在本地數(shù)據(jù)集上計算梯度后,執(zhí)行本地訓(xùn)練后并上傳模型參數(shù),服務(wù)器對各個參與方的模型進(jìn)行加權(quán)平均,從而實現(xiàn)一次系統(tǒng)層面的迭代.但是在FedAvg中,只考慮各數(shù)據(jù)源的權(quán)重與數(shù)據(jù)源的數(shù)據(jù)量大小,并未考慮數(shù)據(jù)源之間數(shù)據(jù)質(zhì)量的差異,若出現(xiàn)數(shù)據(jù)源的數(shù)據(jù)量很大,但其數(shù)據(jù)質(zhì)量很差,這就會造成權(quán)重的不合理分配.因此Huang 等[17]在FedAVG基礎(chǔ)上提出自適應(yīng)數(shù)據(jù)增強(qiáng)的LoAdaBoost FedAvg算法,利用了數(shù)據(jù)共享策略與全局損失函數(shù)的中位數(shù),但是該方法由于損失函數(shù)的限制,很難達(dá)到較高的性能.此外,Lin使用本地模型的平均對數(shù)進(jìn)行匯總并且聚合[18],來防止聚合過程中梯度遺忘的問題.這些方法雖然有效地解決了一定的異質(zhì)性給全局模型帶來的影響,但是這些模型的前提是假設(shè)所有的參與方是誠實,如果存在惡意的參與方或者極端數(shù)據(jù)時,很容易引起全局模型過程中梯度爆炸的問題,并且使得全局模型低于原本參與方模型的性能.

2.3 聯(lián)邦蒸餾算法

本文中,解決通信開銷問題的主要工具是知識蒸餾模型.它是Hinton提出的一種師生結(jié)構(gòu),用于在深度學(xué)習(xí)間傳遞知識.盡管框架最初只是為了在同一數(shù)據(jù)集上的邊緣轉(zhuǎn)移,但是之后,通過在學(xué)生網(wǎng)絡(luò)的訓(xùn)練過程中,引入了一種正則化,挖掘兩個模型傳遞過程中的潛力.它主要思想是將大型神經(jīng)網(wǎng)絡(luò)學(xué)習(xí)到的知識提煉并傳授給小模型中,從而實現(xiàn)從大型網(wǎng)絡(luò)壓縮成小型網(wǎng)絡(luò)的目的.Jeong等學(xué)者[19]就提出了聯(lián)邦蒸餾(Federated Disillation,FedDisill)相對較實用且效果較好的一種方法,它能夠大大減少通信代價的同時適用于缺少標(biāo)簽的異質(zhì)數(shù)據(jù)、異構(gòu)模型的場景,基本概念就是只交換本地局部模型的輸出而非交換傳統(tǒng)聯(lián)邦學(xué)習(xí)采用的全局模型參數(shù),這使得輸出的模型參數(shù)尺寸比傳統(tǒng)的小得多,從而可以減少通信開銷.隨后,就有Gihun Lee[20]學(xué)者認(rèn)為全局模型的收斂性受到數(shù)據(jù)異質(zhì)性的影響,全局模型會對有偏差的客戶端的產(chǎn)生梯度偏移,于是提出了聯(lián)邦本地自我蒸餾(Federated Local Self-Distillation,FedLSD)模型用來緩解梯度偏移的問題,并提高其性能.但是聯(lián)邦蒸餾的不足之處在于只能緩解異質(zhì)性的數(shù)據(jù)產(chǎn)生梯度遺忘的問題,并不能解決出現(xiàn)惡意或極端模型下的梯度爆炸的問題,并且當(dāng)兩個網(wǎng)絡(luò)模型大小相差太大時,知識蒸餾的效果會失效,因此聯(lián)邦蒸餾的收斂性和應(yīng)用性研究需要進(jìn)一步研究.

3 U-FedDO方法

基于FedAvg算法在non-IID的數(shù)據(jù)集導(dǎo)致性能下降的問題,本文提出了一種在服務(wù)器篩選的機(jī)制,并應(yīng)用在教師學(xué)生模型中.本節(jié)分別討論性能下降對應(yīng)的數(shù)學(xué)模型與調(diào)節(jié)器上的約束條件,并且介紹調(diào)節(jié)器的篩選方式、參與方中Unet教師學(xué)生結(jié)構(gòu)的變換、U-FedDO算法的工作流程.

3.1 數(shù)學(xué)模型

FedAvg是由隨機(jī)梯度下降算法(stochastic gradient descent,SGD)推演而來的,假設(shè)有K客戶參加聯(lián)邦學(xué)習(xí),則每個參與方在本地數(shù)據(jù)上使用SGD計算梯度,將模型參數(shù)上傳給服務(wù)器,然后服務(wù)器將收到模型參數(shù)加權(quán)平均并分配給參與方,經(jīng)過一定次數(shù)的迭代后停止的.各參與方傳遞的梯度下降算法的優(yōu)化目標(biāo)為:

(1)

(2)

本文中,聯(lián)邦蒸餾的核心思想是客戶端在本地的自我蒸餾,若全局模型聚合過程中對有偏見的局部分布進(jìn)行擬合,對應(yīng)的權(quán)重將發(fā)生偏移,這很容易導(dǎo)致性能的降低.因此,可以通過蒸餾的方法保留關(guān)于本地數(shù)據(jù)的分布,當(dāng)全局模型聚合導(dǎo)致偏移時,本地的教師模型可以糾正全局模型的分布.以下?lián)p失函數(shù)LFed DO是交叉熵LCE和蒸餾損失LLSD的線性組合:

LFed DO=(1-β)LCE+βLLSD

(3)

其中,q(x,w)是每個本地客戶模型和全局模型在最后一輪的softmax概率與單次真實標(biāo)簽py(x).LLSD是教師模型預(yù)測與學(xué)生模型預(yù)測之間的KL-Divergence蒸餾損失,具體公式為(4):

(4)

其中c∈[C]類的預(yù)測概率,C是分類的數(shù)量.蒸餾的損失是將對數(shù)概率softmax函數(shù)z除以τ來軟化其預(yù)測概率.具體如公式:

(5)

聯(lián)邦學(xué)習(xí)服務(wù)器的權(quán)重根據(jù)FedAvg算法計算.因此結(jié)合公式(2)、公式(3)可以得出,聯(lián)邦蒸餾模型的更新參數(shù)公式(6):

(6)

(7)

(8)

D[f(x,wt+1),y]>D[f(x;w(k)),y]-d

(9)

其中約束條件(8)保證每次篩選中聚合的客戶端參數(shù)參與聚合并且所有的權(quán)重和為1,約束條件(9)中的D為評價指標(biāo),d為評價指標(biāo)的訓(xùn)練過程中的相對誤差值.

3.2 算法模型

本次設(shè)計的U-FedDO算法主要有服務(wù)器、調(diào)節(jié)器、參與方的學(xué)生模型與參與方的教師模型.服務(wù)器負(fù)責(zé)全局權(quán)重的聚合與輸出給本地模型,調(diào)節(jié)器是本地傳遞的權(quán)重進(jìn)行篩選,參與方中有教師模型保留本地模型的權(quán)重,學(xué)生模型主要負(fù)責(zé)訓(xùn)練服務(wù)器傳遞的權(quán)重與結(jié)合本地教師模型的損失.

客戶端中需要將Unet的模型壓縮并參與全局訓(xùn)練,但傳統(tǒng)的蒸餾法并不足以將Unet模型中的U型結(jié)構(gòu)蒸餾成一個非常小的U型網(wǎng).因此在學(xué)生模型的設(shè)計過程中,先是遵循了經(jīng)典Unet的編碼器-解碼器的結(jié)構(gòu)、編碼器為4個下采樣包括3×3的卷積和2×2的池化層、單通道圖像作為輸入,但不同的是在將傳統(tǒng)Unet的卷積層數(shù)從64改成16,同時保持收縮路徑翻倍的趨勢,具體如表1 所示.上采樣為解碼器也是通過4個2×2的池化層,同時保持傳統(tǒng)Unet在每個池化層后面采用Relu激活函數(shù)的結(jié)構(gòu),來防止梯度消失的問題,與降低計算量的問題.最后一層采用softmax激活函數(shù)對像素進(jìn)行分類.在醫(yī)學(xué)圖像分割任務(wù)中,輸入圖像與輸出掩碼在外觀表現(xiàn)上是有差異的,但在底層結(jié)構(gòu)和特征上是一致的.在連續(xù)的卷積過程中會導(dǎo)致某些潛在特征被忽略,因此在特征映射掩碼過程中,在下采樣和上采樣之間加入跳躍連接可以保持高級特征并映射到分割掩碼.在醫(yī)學(xué)圖像分割中,通常只有病灶與背景兩種分類,損失函數(shù)則為二元交叉熵(Binary Cross Entropy,BCE),其公式為(10):

表1 Unet卷積層數(shù)對比表Table 1 Unet convolutional layer comparison table

(10)

其中,N為輸入圖像像素總數(shù),gi是第i個像素真實的標(biāo)簽,0表示背景,1為肺部.pi是對應(yīng)像素預(yù)測為前景的概率.當(dāng)前景像素遠(yuǎn)遠(yuǎn)小于背景像素數(shù)量時,即gi=0的數(shù)量遠(yuǎn)大于gi=1的數(shù)量,在優(yōu)化過程中g(shù)i=0的部分便會占據(jù)主導(dǎo),使得模型嚴(yán)重偏向背景.因此,在客戶端中主要有兩個模型進(jìn)行訓(xùn)練,第一是傳統(tǒng)的64-Unet模型作為教師模型在本地訓(xùn)練,第二是改良后的學(xué)生16-Unet模型參與全局訓(xùn)練,通過公式(4)的蒸餾損失函數(shù),學(xué)生模型能夠繼承教師模型訓(xùn)練概率分布,保證性能不會因卷積層的降低而下降.

服務(wù)器主要功能是將被選擇的參與方上傳的本地模型進(jìn)行聚合,并將重復(fù)發(fā)放給各個參與方.服務(wù)器中就需要三個階段分別為初始化、聚合、評估.在初始化階段,服務(wù)器首先將醫(yī)療機(jī)構(gòu)加入聯(lián)邦學(xué)習(xí)模型訓(xùn)練任務(wù)中,同時服務(wù)器需要分配信息參數(shù)發(fā)送給參與方如學(xué)習(xí)率、迭代次數(shù)、優(yōu)化器等.在聚合階段,服務(wù)器需要接受客戶端上傳的模型參數(shù),并使用將收到的模型參數(shù)進(jìn)行聚合,聚合公式為(11):

(11)

調(diào)節(jié)器主要功能是篩選與重新平衡.首先,通過服務(wù)器設(shè)置的約束機(jī)制判斷對各個參與方傳遞的模型參數(shù),若為有利方則參加聚合,若為有害方則剔除.同時調(diào)節(jié)器還根據(jù)參與方的數(shù)量重新分配模型的權(quán)重,若參與方多次被篩選出去時,調(diào)節(jié)器將不再參加后續(xù)模型的迭代,從而減少通信開銷.此外,在公式(9)的篩選給定相對誤差值d,緩解模型聚合偏移帶來的局部收斂的情況.

3.3 工作流程

如圖1,顯示U-FedDO框架的訓(xùn)練、傳遞、蒸餾、篩選、聚合的過程.首先,服務(wù)器進(jìn)行初始化設(shè)定模型參數(shù),學(xué)習(xí)率迭代次數(shù)等,接著客戶端訓(xùn)練一個教師64-Unet模型,同時接受服務(wù)器傳來的初始化全局16-Unet模型的參數(shù),開始進(jìn)行訓(xùn)練.然后,客戶端將訓(xùn)練好的參數(shù)發(fā)送給服務(wù)端,服務(wù)端中通過調(diào)解器去篩選能夠提高準(zhǔn)確度的參與方,并更新每個參與方在服務(wù)器聚合的權(quán)重.將聚合后的模型參數(shù)再依次分配給參與方進(jìn)行下輪訓(xùn)練,直到模型收斂,保存模型并結(jié)束本次的聯(lián)邦學(xué)習(xí).此外,調(diào)解器是虛擬組件,它可以直接部署在聯(lián)邦服務(wù)器上以減少通信開銷.具體如算法1所示.

圖1 U-FedDO結(jié)構(gòu)圖Fig.1 U-FedDO structure diagram

算法1.基于Unet的聯(lián)邦蒸餾學(xué)習(xí)

輸入:全局權(quán)重,本地迭代次數(shù),本地數(shù)據(jù),學(xué)習(xí)率

1.在服務(wù)器執(zhí)行

2.初始化原模型參數(shù)W0,將原始的模型參數(shù)W0,廣播給所有的參與方

3.For 每一全局模型更新輪次t=1,2,…,do

4. 調(diào)節(jié)器通過篩選確定C個參與方的集合

5. For 所有參與方C并行地 do

8. End for

10.If 模型收斂

12.End if

13.在調(diào)節(jié)器執(zhí)行執(zhí)行

14.從參與方獲得最新的本地模型參數(shù)與服務(wù)器上的測試集

15.For 所有參與方C

16. ifD[f(x;wt+1),y]>D[f(x;w(k)),y]-d

17.將可以用模型參數(shù)存入?yún)?shù)Ct

18. End if

19.確定參與存入的參數(shù)與數(shù)量Ct,并發(fā)生給服務(wù)器

20.End for

21.在參與方更新

24.for 從1到迭代次數(shù)S的每一個本地迭代I do

25. for 從1到批量數(shù)量B=的批量序號b do

28. End for

29.End for

4 實驗及結(jié)果分析

4.1 實驗數(shù)據(jù)與配置參數(shù)

本次分別使用COVID-19[21]掃描影像與HAM10000[22]數(shù)據(jù)集.其中,COVID-19掃描影像中有20張已標(biāo)記的圖像,共切片為3530張圖片,左右肺與感染由兩名放射科醫(yī)生標(biāo)記,并由經(jīng)驗豐富的放射科醫(yī)生進(jìn)行驗證;HAM10000數(shù)據(jù)集包含10000張七種互斥類的皮損區(qū)域皮膚鏡圖像,且至少在3位皮膚科專家共識的基礎(chǔ)上對其標(biāo)記,并通過金標(biāo)準(zhǔn)組織病理學(xué)的證實.

本文采用留出法策略,即分別在K-1源域訓(xùn)練,并在一個未見過的目標(biāo)域上進(jìn)行測試.因此,將COVID-19掃描影像與HAM10000數(shù)據(jù)劃分為本地訓(xùn)練集,本地驗證集、本地測試集、全局驗證集、全局測試集.在IID數(shù)據(jù)集的實驗設(shè)計中,將兩個數(shù)據(jù)集分別以10%的全局驗證、10%的全局測試,并將剩下的80%的數(shù)據(jù)隨機(jī)均分至3個醫(yī)療機(jī)構(gòu),并且機(jī)構(gòu)也將分配的數(shù)據(jù)以8∶1∶1的比例分為本地訓(xùn)練、本地驗證、本地測試.由于Unet的圖像分割模型本身可適用于小數(shù)據(jù)集,且U-FedDO就是解決醫(yī)院數(shù)據(jù)不足與數(shù)據(jù)不平衡的問題,則在non-IID數(shù)據(jù)的實驗設(shè)計中,將兩個數(shù)據(jù)集中隨機(jī)抽取以4∶2∶2∶1的比例的小數(shù)據(jù)存放在4個不同的機(jī)構(gòu),其余數(shù)據(jù)與IID數(shù)據(jù)集一致.此外,所有的訓(xùn)練集與測試集之間是互斥的關(guān)系.

實驗硬件配置為:CPU8核,內(nèi)存28G,顯卡 RTX 3080,軟件系統(tǒng)包括Ubuntu21.0,CUDA8.0,本次以單機(jī)的模式進(jìn)行評估.在聯(lián)邦學(xué)習(xí)過程中,所有客戶使用相同的超參數(shù)設(shè)置,本地模型使用GDB優(yōu)化器進(jìn)行訓(xùn)練,批量大小為1,動量為0.9.步驟大小和學(xué)習(xí)速率都被設(shè)置為0.01,圖片壓縮為0.5尺寸,知識蒸餾溫度參數(shù)τ根據(jù)經(jīng)驗設(shè)定為10,共訓(xùn)練了50個全局輪次,每個全局輪次中的局部歷時迭代設(shè)置為1.

4.2 評價指數(shù)

為了對分割性能進(jìn)行評測,本文采用的評價指標(biāo)包括Dice相似系數(shù)來衡量分割模型的性能,這是一個用于醫(yī)學(xué)圖像分割的標(biāo)準(zhǔn)評價指標(biāo).具體的計算公式如下:

(12)

其中Sseg為預(yù)測值,Ssef為真實值.

4.3 消融研究

“消融研究”在機(jī)器學(xué)習(xí)中經(jīng)常被用一些來描述去除網(wǎng)絡(luò)的某些部分的過程,以便更好地理解網(wǎng)絡(luò)的行為.U-FedDO模型中的3個關(guān)鍵問題:1)教師-學(xué)生的16Unet模型對性能與通信開銷的影響情況;2)U-FedDO在數(shù)據(jù)平衡與非平衡下的性能情況;3)參與的醫(yī)療機(jī)構(gòu)數(shù)量增加情況下,U-FedDO性能變化.

1)在模型參數(shù)方面,如表2所示傳統(tǒng)的深度神經(jīng)網(wǎng)絡(luò)模型通常有數(shù)百萬,數(shù)千萬甚至上億的參數(shù)量,因此發(fā)送如此多的浮點數(shù)值將給協(xié)調(diào)方產(chǎn)生巨大的通信開銷,而本文設(shè)計的輕量16-Unet模型明顯不管是模型大小還是參數(shù)數(shù)量都遠(yuǎn)小于其他模型,同時也比傳統(tǒng)的Unet模型減少了17倍的參數(shù).訓(xùn)練時間方面,在512×512像素的數(shù)量為1000的切片圖片中,傳統(tǒng)的64-Unet模型單次訓(xùn)練時間為84s,而16-Unet模型的單次訓(xùn)練時間為63s,訓(xùn)練時間降低約為25%.

表2 經(jīng)典CNN模型與Unet的模型大小與參數(shù)數(shù)量Table 2 Model size and number of parameters of classical CNN model and Unet

在性能方面,如圖2所示,在64-Unet為經(jīng)典的Unet模型,而16-Unet為經(jīng)過知識蒸餾后的模型,可以判斷的出模型在局部驗證普遍比在全局驗證中的準(zhǔn)確率更高,而輕量的16-Unet在教師-學(xué)生模型下,性能不會大量的下降.

圖2 16-Unet與64-Unet在測試集中的性能對比圖Fig.2 Performance of 16-Unet vs.64-Unet in the test set

2)U-FedDO的分割效果主要在兩個方面分別是IID數(shù)據(jù)與non-IID數(shù)據(jù).一是在IID數(shù)據(jù)集下,如圖3所示,兩個數(shù)據(jù)集在FedAvg、U-FedDO、FedDistill中的差距不大,但是由于FedDistill中的本地教師模型會保留本地的性能較差的模型參數(shù)與輕量Unet的模型的影響下,性能會有所下降,但是U-FedDO在有效地篩選情況下,性能與FedAvg幾乎一致.同時對比圖2中單個機(jī)構(gòu)的性能可以得到以下結(jié)論,當(dāng)用單源數(shù)據(jù)訓(xùn)練的模型在部署到未知的領(lǐng)域時不能獲得良好的結(jié)果,但是在更多的客戶參與的聯(lián)邦訓(xùn)練時,泛化性能會增加,這是由于從多個來源聚集的數(shù)據(jù)可以覆蓋一個更全面的數(shù)據(jù)分布.

圖3 3種方法在IID數(shù)據(jù)下的性能對比圖Fig.3 Performance of 16-Unet vs.64-Unet in the test set

二是在non-IID數(shù)據(jù)下,如圖4所示,4個醫(yī)療機(jī)構(gòu)的準(zhǔn)確率也是各不相同,本文設(shè)計的U-FedDO算法在兩個數(shù)據(jù)集中的non-IID數(shù)據(jù)性能是最高的,最高機(jī)構(gòu)的性能還高約10%,且遠(yuǎn)比FedAvg的準(zhǔn)確度高,這取決于調(diào)節(jié)器在聚合過程中,有效地篩選出其他惡意的參與方模型,保證在服務(wù)器在聚合過程始終對全局模型的優(yōu)化的效果.

圖4 4個醫(yī)療機(jī)構(gòu)與3個聯(lián)邦模型的分割箱型圖Fig.4 Split box plot of four medical institutions with three federal models

為了展現(xiàn)顯示效果,圖5顯示了COVID-19掃描影像與HAM1000的部分測試集的分割效果,可以看到在本次實驗的數(shù)據(jù)集中,U-FedDO分割的效果在整體的形狀和位置上都已經(jīng)十分接近手工標(biāo)注,對于局部的圖像信息也展現(xiàn)得較為清晰與詳細(xì),線條相對逼真與準(zhǔn)確,但是在邊界上手工標(biāo)注的會比U-FedDO更圓滑一點.而相比與FedAvg與FedDistill兩個算法由于極端數(shù)據(jù)的影響,在COVID-19掃描影像雖然整體的肺部結(jié)構(gòu)大致能分割出來,但是在肺部外無相關(guān)的也進(jìn)行了誤判分割,顯得分割效果有雜亂且沒有目標(biāo);在HAM1000的數(shù)據(jù)集中,FedAvg與FedDistill分割效果也低于U-FedDO,特別是當(dāng)圖像中還有毛發(fā)干擾時,分割效果的差距尤為明顯.

圖5 3種方法不在平衡數(shù)據(jù)下的分割對比圖Fig.5 Comparison of the segmentation of the three methods not under the balanced data

3)當(dāng)不同醫(yī)院數(shù)量且醫(yī)院間的數(shù)據(jù)分布不平衡的情況下,U-FedDO和FedAvg的性能會有怎么樣的變化.圖6展示了圖像分割的結(jié)果,參與聯(lián)邦訓(xùn)練的醫(yī)療機(jī)構(gòu)依次從1逐漸增加到16,并且圖片數(shù)量不同且每個參與方的準(zhǔn)確率也各不相同.U-FedDO模型在1至16中的準(zhǔn)確率始終高于傳統(tǒng)的FedAvg,并且隨著參與方數(shù)量的變化趨于平穩(wěn)上升.FedAvg由于存在大量的數(shù)據(jù)不平衡的參與方所導(dǎo)致的準(zhǔn)確率不高的情況,特別是當(dāng)數(shù)量到9與10時,參與方存在極度不平衡狀態(tài)且準(zhǔn)確率不高的情況,由此可以判斷使用FedAvg會明顯產(chǎn)生梯度爆炸的問題.本次實驗再次證明了U-FedDO在數(shù)據(jù)不平衡情況下,能夠有效地防止梯度爆炸的情況,并且泛化能力方面具有穩(wěn)定的功效.

圖6 醫(yī)療機(jī)構(gòu)不同數(shù)量性能對比圖Fig.6 Comparison of different numbers of medical institutions

5 結(jié)束語

本文中,基于在醫(yī)療圖像分割在聯(lián)邦模型中兩大問題:通信開銷和數(shù)據(jù)不平衡,設(shè)計一個U-FedDO算法,主要是應(yīng)用了模型的知識蒸餾與服務(wù)器上部署調(diào)解器的兩個方法,使用知識蒸餾壓縮方法在經(jīng)典圖像分割模型Unet的結(jié)構(gòu)特性減少其卷積層數(shù),設(shè)計一個輕量的Unet模型.在調(diào)解器方面,通過約束條件與重新設(shè)計聚合權(quán)重的方法保持全局模型不會因為極端數(shù)據(jù)而梯度爆炸.實驗結(jié)果表明,在使用改進(jìn)后的Unet蒸餾方法,模型以極小的性能下降的代價,有效將單次訓(xùn)練的模型參數(shù)降低了16倍并且訓(xùn)練時間減少約25%,大量的降低了聯(lián)邦訓(xùn)練過程的總時耗與通信數(shù)據(jù)總量.在IID數(shù)據(jù)集的實驗中,FedAvg、U-FedDO、FedDistill性能并沒有很大差異.在non-IID的數(shù)據(jù)中,U-FedDO模型明顯比FedAvg的準(zhǔn)確更高,且隨著參與方的數(shù)量增加,U-FedDO也比FedAvg得性能表現(xiàn)的更穩(wěn)定.

研究發(fā)現(xiàn),本文只針對橫向聯(lián)邦學(xué)習(xí)的不均衡數(shù)據(jù)做探討,在下一階段的工作中,將深入研究縱向聯(lián)邦學(xué)習(xí)在圖像分割的應(yīng)用.