BT - Unet:生物醫學圖像分割的自監督學習框架
BT-Unet采用Barlow twin方法對U-Net模型的編碼器進行無監督的預訓練減少冗余信息,以學習數據表示。之后,對完整網絡進行微調以執行實際的分割。
BT-Unet由 Indian Institute of Information Technology Allahabad開發,發布在2022年的JML上。
BT-Unet
BT-Unet架構圖:a、預訓練U-Net編碼器網絡,b、用預訓練的編碼器權值初始化的微調U-Net模型
BT-Unet框架分為兩個階段:1)預訓練階段和2)微調階段。
預訓練
預訓練的目的是使用無注釋的數據樣本,學習復雜的特征表示。U-Net模型的編碼器使用Barlow Twins (BT)策略進行預先訓練,然后進行微調以執行實際的分割:
BT-Unet框架可應用于各種先進的U-Net模型:經典U-Net、注意力U-Net (A-Unet)、inception U-Net (I-Unet)和residual cross-spatial attention guided inception U-Net (RCA-IUnet)。
微調
U-Net 模型中編碼器網絡的權重使用預訓練權重(來自第一階段)進行初始化,而網絡的其余部分使用默認權重進行初始化。
使用有限的注釋樣本對 U-Net 模型進行微調,用于生物醫學圖像分割。
U-Net 模型使用分段損失函數進行微調,L 定義為二元交叉熵損失、LBC 和dice coefficient損失的平均值,LDC:
其中,y為像素的真值標簽,p(y)為像素的預測標簽,N為像素的總數量。
結果表現
論文使用下面的數據集進行評測對比:
基于相同樣本數量的小訓練集的性能分析
- KDSB18:BTU-Net 模型的性能優于沒有使用 BT方法 的模型。
- BUSIS:U-Net 和 A-Unet 模型無法學習和提取有關腫瘤區域的特征圖(精度、DC 和 mIoU 為 0),但是通過預訓練,這些模型取得了顯著的改進。在 I-Unet 和 RCAIUnet 模型的情況下,通過預訓練可以得到相當大的改進。
- ISIC18:I-Unet 和 RCAIUnet 模型是影響最大的網絡,精度分別提高了 5.1% 和 2.2%。然而,在使用 BT 預訓練時,經典的U-Net 和 A-Unet 的性能略有下降。
- BraTS18:I-Unet 和 RCA-IUnet 模型在使用 BT-Unet 框架時在分割性能上取得了顯著提升,而普通 U-Net 和 A-Unet 模型則沒有觀察到相同的行為。
在不同規模小型訓練集的性能研究
對于所有訓練數據比例小于50%的數據集,在模型之間觀察到類似的性能變化。
定性結果
帶有BT的RCA-IUNet具有很好的分割效果。