基于Unet 的多任務醫學圖像語義分割模型*

沈旭東 ，樓平* ，吳湘蓮 ，朱立妙 ，雷英棟

(1.嘉興職業技術學院智能制造學院,浙江 嘉興 314036；2.同濟大學浙江學院機械與汽車系,浙江 嘉興 314051)

氣胸是指氣體進入胸膜腔,造成積氣狀態,氣胸對于人類的生命是很大的危險,一般判斷正確率僅在50%,其中通常的方法是通過X 光片進行診斷,但是X 光片須有經驗豐富的醫生才能判斷。當前隨著深度學習技術[1]的不斷發展,通過深度學習模型就可以很快地識別出是否患有氣胸疾病,并且能夠判斷出病灶的位置,而且通過網絡模型預測病人是否患有氣胸也將很大程度上減輕醫生的負擔,可以提供有效的輔助手段。

在使用傳統Unet 模型[2-3]進行語義分割的過程中,遇到的第一個問題是,正樣本召回率低,如圖1 所示,X 光片有氣胸病灶,但是通過語義分割模型預測的過程,沒有識別出來；第二個問題是,能否在原有的基礎上,設計更好的網絡結構,進一步提高語義分割的精度。第一個問題原因在于Unet 是針對點對點的像素級別的語義分割模型,不能從全局的角度進行有無病灶的判斷,因此對于小的病灶,單純地使用Unet網絡進行分割,無法得到很好的預測結果,本文在經典Unet 網絡的基礎上做了改進,引入“分類”和“分割”多任務策略進行學習,通過對不同的圖像在全局上進行有無病灶二分類學習,再對每個圖像上的單個像素進行語義分割,將兩者的損失函數進行融合,通過這種方法可以提高正樣本預測的準確率。對于第二個問題,傳統的解決方式是增加網絡的體積,但是調大神經元的數量來設計更深層次的網絡結構,將會導致占用非常大的硬件資源,而且也會使得訓練變得非常困難[4],本文在網絡結構的改進上,采用Resnet34[5-6]作為編碼框架,設計了空間通道壓縮與激發(Spatial-Channel Sequeeze ＆ Excitation,SCSE)模塊進行修正,提高訓練預測的精度。

圖1 氣胸CT 圖像預測錯誤

1 網絡模型結構設計

1.1 總體設計

本文設計了一種針對氣胸CT 圖像的端到端語義分割模型,采用編碼、解碼的網絡結構,在設計二分類網絡和語義分割網絡時,其中編碼階段的基礎特征提取網絡是很重要的一部分,由于獲取海量氣胸CT 圖像的難度比較大,本文采用以Resnet34 為框架的帶預訓練參數的Unet 模型作為特征提取網絡,因為它通過在網絡中增加殘差網絡的方法,解決了網絡深度到一定程度,誤差升高,效果變差,梯度消失現象明顯[7],使得網絡反向傳播求最小損失難以實現的問題。在編碼階段,引入SCSE 模塊結構,得到圖像二分類損失,解碼時融合編碼階段各層特征信息,進行上采樣,對上采樣的各層特征信息再進行融合,得到語義分割圖像,再對圖像二分類和圖像語義分割兩者的損失函數進行融合,調整權重關系,得到全局損失函數,再進行優化,最終得到預測圖像。具體結構如圖2 所示。

圖2 網絡模型結構設計

1.2 SCSE 模塊設計

假設輸入圖像X∈?H×W×C,通過編碼模塊Resnet34 和解碼模塊上采樣生成輸出特征圖U∈?H×W×C′,H和W表示圖像的空間高度和寬度,C和C′表示輸入和輸出的通道數量,通過卷積和非線性變換操作后,U中包含了X中空間和通道信息,我們在每次編碼或解碼操作后,引入SCSE 模塊,將U修正為U′,SCSE 模塊實現方式如圖3 所示,主要有空間壓縮模塊和通道壓縮模塊兩部分組成。

圖3 SCSE 模塊結構圖

空間壓縮模塊設計如下,假設輸入特征圖U＝[u1,u2,…uC],其中ui∈?H×W代表每個通道的特征圖,通過全局平均池化[8-9](global average pooling,GAP)操作,得到向量z∈?1×1×C,對于第k個通道特征圖變換如式(1):

該變換在z中引入了全局空間特征信息,再通過z′＝W1(δ(W2z)),得到z′,其中W1∈,W2∈為兩個全連接網絡,δ(·)為ReLU[10]操作。再通過sigmoid 變換,得到歸一化修正系數σ(z′)∈[0,1],最終通過修正系數乘以輸入特征圖U,得到,如式(2):

通過該模塊,不重要的通道內信息會減小,被抑制,而重要通道內的信息幾乎保持不變,變相得到增強。

通過上述空間壓縮信息,從而獲得衡量通道重要性的指標,我們也可以通過壓縮通道信息,獲取衡量空間位置重要性的指標。同樣對于輸入特征圖U＝[u1,u2,…uC],通過轉換得到

式中:ui,j∈?1×1×C代表對應的空間位置(i,j),通過q＝Wsq*U卷積操作實現空間擠壓操作,再通過sigmoid變換,得到歸一化修正系數σ(q)∈[0,1],最終通過修正系數乘以輸入特征圖U,得到,如式(4):

通過相加得到兩者的融合信息,最后得到對空間信息和通道信息重要性都做了權重處理的輸出。如式(5):

1.3 多任務學習策略

在對氣胸CT 圖像進行語義分割時,為了減少噪聲干擾,當分割得到的圖像像素點求和小于一定數量時,預測該圖像為無病灶特征,因此相當于用語義分割算法做了二分類操作,但是語義分割模型以及損失函數只考慮了單個像素點的分類,沒有考慮整張圖片所有像素點求和的整體分類[11-12],因此在優化的過程中只是針對局部進行優化,沒有考慮到對整體進行優化,針對這個問題,本文引入了多任務學習策略,直接使分類和語義分割兩個網絡共用同一個網絡結構,提取特征圖,而不是分別使用兩個獨立的網絡,結構如圖4 所示,在保證整個網絡結構不變的情況下,設置多個任務“分類”和“分割”,對于二分類,本文引入邏輯斯特損失函數loss1,對于分割任務,本文引入交叉熵損失函數loss2,在訓練時,首先通過單獨的兩個損失函數分別計算分類和分割兩個輸出的損失,再計算這兩個損失的加權和作為網絡的損失,最終得到融合的損失函數loss＝a×loss1＋b×loss2,其中a、b為兩者的權重值。

圖4 多任務學習結構圖

2 實驗和分析方法

2.1 數據集及處理方法

本文所用數據集包含了原始圖像及標注的分割圖像,標注的分割圖像由RLE(Run Length Encoding)編碼,整個數據集包含12 047 張原始圖像及標注好的分割圖像,如圖5 所示,尺寸為1 024×1 024 pixel,將所有圖像分為訓練集數據和驗證集數據,其中訓練集占90%,驗證集占10%。由于GPU 容量的限制,將原始1 024×1 024 pixel 圖像重新整形為512×512 pixel,為了進一步擴大數據集的容量和提高網絡模型的泛化能力,將圖片進行隨機水平翻轉、垂直翻轉、旋轉、亮度變換,高斯模糊,透視變化等操作,隨機比例為50%。

圖5 帶標注的數據集

2.2 評價指標

評價指標用于衡量網絡模型的好壞,可以說明模型的性能,辨別模型的結果。本文所采用的模型使用Dice 作為評價指標[13],Dice 系數是一種集合相似度度量指標,通常用于計算兩個樣本的相似度,取值范圍為[0,1],分割結果最好時值為1,最差時值為0,在該模型中,用于檢測預測病灶和真實病灶的相似度,如式(6):

2.3 訓練相關細節

本文在設計相關實驗時,采用pytorch 框架進行設計,實驗所用圖形工作站配置為:4 核CPU,內存16G,GeForceGTX1080Ti GPU,顯存16G,操作系統Ubuntu 16.04,網絡模型配置如下:優化器采用隨機梯度下降,momentum＝0.9,weight_decay＝0.0001,LearningRate＝0.002,由于單張圖片為512×512 pixel,圖片較大,批處理大小設置為4,為了增加批處理大小,本文采用梯度累加實現“顯存擴大”,進行N次前向后反向更新一次參數,相當于擴大了N倍的batchsize,即每次獲取1 個批處理的數據,通過梯度下降法,計算梯度,梯度不清空,不斷累加,累加一定次數后,根據累加的梯度更新網絡參數,然后清空梯度,進行下一次循環。一定條件下,batchsize 越大訓練效果越好,梯度累加則實現了批處理量的變相擴大,本文設置累加次數為4 倍,相當于將批處理大小設置為16。

2.4 實驗結果及分析

(1)多任務學習策略對比分析

本文在多任務學習策略中,采用對“分類”和“分割”兩種損失函數進行加權平均的方式,根據多任務損失函數loss＝a×loss1＋b×loss2,a、b為兩者的權重值,對a、b兩者取不同的值進行對比分析,實驗結果如表1 所示,a、b實際為氣胸圖像良惡性判別和語義分割的損失權重之比,在訓練時,有可能出現某一任務的偏重明顯大于其他任務的情況,這樣可能使有些任務的訓練效果較差,所以在這里加入損失權重,來對每個任務的比重進行調節,以防止某個任務的訓練效果較差。針對不同的權重系數,采用Unet＋Resnet34＋SCSE 網絡模型進行對比分析,實驗表明對于該任務,分類loss 權重a＝0.1,語義分割loss 權重b＝0.9 時預測效果最好。

表1 不同權重系數預測對比結果分析

(2)不同網絡模型對比分析

不同網絡模型對比分析結果如表2 所示,本文設計的模型在預測氣胸CT 圖像時,其中Dice 相似系數得分達到85.3%,成績與目前最好“五折交叉驗證訓練模型”接近,但簡化了模型結構和訓練復雜度,實驗表明采用Resnet34 作為編碼框架,引入SCSE 模塊將比單Unet 網絡提高約4%,采用二分類和語義分割雙任務進行預測比單任務進行預測在預測精度上也提高1.2%。由此可以分析,引入SCSE模塊對圖像特征信息進行修正,從空間和通道兩個方向獲取圖像的全局信息,比傳統的Unet＋Resnet34模型對于特征的獲取更加有效,損失函數采用“分類”和“分割”融合的多任務策略對氣胸醫學圖像進行語義分割,可以更好地優化訓練過程,避免了只是針對局部優化,沒有考慮到對整體進行優化的缺點。

表2 不同網絡模型對比結果分析

(3)預測效果

根據訓練得到的最優網絡模型Unet＋Resnet34＋SCSE＋Multi-Task 組合,對實際醫學圖像進行分割,以3 個病人為例進行預測,預測效果如圖6 所示。

圖6 實驗預測效果

3 結束語

深度學習模型是醫學圖像檢測領域非常強大的工具,使用語義分割預測氣胸CT 圖像病灶位置對醫生判別疾病將起到很好的輔助作用,本文設計的語義分割模型采用Resnet34 作為框架,解決了網絡深度到一定程度,梯度消失的問題；引入SCSE 模塊,從空間和通道兩個方向獲取圖像的信息,提高了語義分割的精度；損失函數采用“分類”和“分割”多任務策略進行學習,從局部像素點和全局圖像的分類上進行融合,進一步提高了語義分割的精度。