如何利用深度學習診斷心臟病
人類心臟是一臺令人驚嘆的機器,它能持續運轉長達一個世紀而不失靈。測量心臟功能的關鍵方法之一是計算其 射血分數 ,即每搏輸出量占心室舒張末期容積量的百分比。而測量這個指標的第一步依賴于對心臟圖像心室的分割。
當我在紐約從事 Insight AI計劃 時,我決定著手處理 右心室分割問題 。下面我將簡單介紹一下整個過程。
問題描述
開發一個能夠對心臟磁共振成像(MRI)數據集圖像中的右心室自動分割的系統。到目前為止,這主要是通過經典的圖像處理方法來處理的。而現代深度學習技術有可能提供更可靠、更自動化的解決方案。
2016年由Kaggle贊助的 左心室分割挑戰賽 中的三名獲獎者都采用了深度學習解決方案。然而,分割右心室(RV)則更具挑戰性,因為:
在腔內存在信號強度類似于心肌的小梁; RV復雜的新月形;分割根尖圖像切片的難度;個體之間的室內形狀和強度存在相當大的差異,特別是在不同的疾病病例之間。
撇開醫學術語不談,要識別RV就更困難了。左心室是一個厚壁圓環,而右心室是一個形狀不規則的物體,有薄的壁,有時會與周圍的組織混合在一起。這是MRI快照右心室內壁和外壁(心內膜和心外膜)的手工繪制輪廓:
這是一個分割起來很容易的例子。這一個比較困難:
而這對于沒有經過訓練的眼睛來說完全是一個挑戰:
事實上,與左心室相比,醫生需要耗費兩倍的時間來確定右心室的體積并生成結果。這項工作的目的是建立一個高度準確的右心室自動分割深度學習模型。模型的輸出是*分割掩碼*,即一個逐像素的掩碼,用來表示某個像素是否是右心室的一部分或只是背景。
數據集
對于當前這個問題,深度學習需要面對的最大挑戰是數據集太小。 數據集(可以訪問 這里 )僅包含了來自于16例患者的243張醫師分割的MRI圖像。 另外還有3697張未標記的圖像,這對于無監督或半監督技術可能會有用,但是我將這些放在了一邊,因為這是一個監督學習問題。 圖像尺寸為216×256像素。
由于數據集過小,人們可能會懷疑無法將其一般化到尚未看到的圖像!但是很不幸,醫療上標記過的數據非常昂貴,并且很難獲取到。對數據集進行處理的標準程序是對圖像應用仿射變換:隨機旋轉、平移、縮放和剪切。此外,我實現了彈性變形,也就是對圖像的局部區域進行拉伸和壓縮。
應用這種圖像增強算法的目的是為了防止神經網絡只記住訓練的樣例,并強迫其學習RV是一個實心的、月牙形的、方向任意的物體。在我實現的訓練框架中,我會隨時對數據集應用圖像變換算法,這樣,神經網絡就會在每次訓練時看到新的隨機變換。
由于大多數像素都屬于背景,所以各個種類之間分布不平衡。如果將像素強度歸一化在0和1之間,那么在整個數據集中,只有5%的像素屬于RV腔的一部分。
在創建損失函數時,我嘗試了重加權方案來平衡種類的分布情況,但最終發現未加權平均算法的表現最好。
在訓練中,有20%的圖像被取出來作為驗證集使用。 RV分割挑戰賽的組織者有一個單獨的測試集,它由另外32個患者的514張MRI圖像組成。我提交的預測輪廓就是使用這個測試集進行最終評估的。
還需要有一種方法來對數據集上的模型性能進行量化。RV分割挑戰賽的組織者選擇使用了 戴斯系數 。模型會輸出一個掩碼*X*來描述RV,而戴斯系數將*X*與由醫師創建的掩碼*Y*通過以下方式進行比較:
計算值是交叉區域與兩區域之和的比率的兩倍。對于不相交的區域,值為0;如果兩區域完全一致,則值為1。
下面我們來看下模型架構。
U-net:基線模型
由于我們只有四周的時間來完成這個項目,所以我想盡快實現一個基線模型并讓它運行起來。我選擇了由Ronneberger、Fischer和Brox提出的u-net模型,因為它曾在生物醫學分割項目中取得過成功,而且它的作者能夠通過使用積極的圖像增強和逐像重新加權算法并僅基于*30張圖片*來訓練網絡。
u-net架構由一個收縮路徑組成,就是將圖像折疊成一組高級特征,隨后是使用特征信息構建像素分割掩碼的擴展路徑。 u-net獨特的地方就是它的“復制和合并”連接,這些連接能夠將信息從早期特征圖傳遞到構建分割掩碼網絡的后續部分。作者指出,這些連接允許網絡同時并入高級特征和像素方面的細節。
我們使用的架構如下所示:
由于我們圖像尺寸是u-net作者原本考慮的一半大小,所以我們需要將原始模型中的降采樣層數從4個減少到3個來適應網格。我們也要用零來填充卷積,以保持圖像的大小不變。該模型是用Keras實現的。
如果沒有圖像增強,u-net在訓練數據集上的戴斯系數能達到0.99(0.01),這意味著該模型具有足夠的能力來捕獲RV分割問題的復雜性。然而,驗證戴斯系數為0.79(0.24),所以u-net過強。圖像增強改進了泛化,并將驗證精度提高到了0.82(0.23),代價是將訓練精度降至0.91(0.06)。
我們如何進一步地降低訓練與驗證之間的差距呢?正如Andrew Ng在這個很棒的 談話 中描述的那樣,我們可以用更多的數據(這不太可能)、正則化(dropout和批處理規范化沒有效果)、或嘗試新的模型架構。
擴張U-net:全局感受野
要對器官圖像進行分割,需要了解器官之間排列的相關知識。事實證明,即使在u-net最深層的神經元也只有68×68像素的感受野。網絡的任何部分都無法“看到”整個圖像。網絡不知道人類只有一個右心室。例如,下面的圖片中,箭頭標記的地方被錯誤分類了:
我們使用擴張卷積來增加網絡的感受野。
在上圖中,底層的卷積是規則的3×3卷積。下一層,我們將卷積擴大了2倍,所以在原始圖像中它們的有效感受野是7×7。如果頂層卷積擴大4倍,則能產生15×15的感受野。以此類推。
從原理上來說,黃色標記的卷積層被u-net中的擴張卷積所替代。最內層的神經元現在具有了覆蓋整個輸入圖像的感受野。我稱之為“擴張u-net”。
在數量上,擴張u-net確實提高了效果,達到了0.85(0.19)的驗證戴斯分數,同時保持了0.92(0.08)的訓練效果!
擴張DenseNet:一次性多個尺度
這個靈感來自于物理中的 張量網絡 ,我決定嘗試使用一個新型的圖像分割架構,我稱之為“擴張DenseNet”。它結合了擴張卷積和DenseNet這兩種想法,這樣能夠大大減少網絡的深度和參數。
對于分割而言,我們需要來自多個尺度的全局上下文和信息來產生像素級掩碼。如果我們完全依賴于擴張卷積來產生全局上下文,而不是通過降采樣來將圖像變得更小呢?現在,所有卷積層的大小都相同,我們可以應用DenseNet架構的關鍵思想,并在所有層之間使用“復制和合并”連接。擴張DenseNet的結果如下圖所示:
在DenseNet中,第一個卷積層的輸出作為輸入饋送到所有的后續層中,第二、第三層也這樣。
擴張DenseNet表現不錯,在驗證集上得到了0.87(0.15)的戴斯得分,訓練精度為0.91(0.10),同時保持了極高的參數效率!
結果
對人類在RV分割方面的評估給如何對模型的表現進行評估指明了方向。研究人員估計,人類完成RV分割任務的戴斯得分為0.90(0.10)。上面所述的已經發布的模型是完全卷積網絡(FCN),測試集上的精度為0.84(0.21)。
我開發的模型在驗證集上已經超過了最新的技術水平,并且正在接近人類的表現!然而,真正的評測是在測試集上評估模型的表現。此外,上面引用的數字是針對心內膜的, 那么心外膜的表現如何呢?我在心外膜上訓練了一個單獨的模型,并將細分輪廓提交給了組織者,希望能獲得最好的成績。
以下是結果,首先是心內膜:
這個是心外膜:
擴張u-net與心內膜上的最新技術水平相當,并超過它在心外膜上的表現。擴張DenseNet緊跟其后,僅有190K個參數。
總結
深度學習模型的表現有時候看起來似乎很神奇,但這是精心設計的結果。即使數據集很小,精心挑選的數據增強方案也可以讓深度學習模型更好地一般化。
根據這些想法,我創建出了最先進的模型來分割心臟MRI中的右心室。我非常高興地看到了擴張DenseNet能夠在其他圖像分割評測上成功運行。
文章的最后,我來談談本項目未來可能需要做的工作:
- 對數據集進行加權,以突出難以分割的根尖切片。
- 探索多級(本地化、注冊、細分)流水線。
- 在生產系統中優化射血分數。
- 內存高效的擴張DenseNet:密集連接的網絡有一個很大的缺點,它需要占用大量的內存。 而TensorFlow的實現卻與眾不同,它將我們限制在16GB GPU并且一個批次具有3個圖像的8個層上。如果切換到之前提出的記憶高效的實現上的話,就可以創建出更深層次的體系架構來。
來自:https://zhuanlan.zhihu.com/p/29465424