天天看點

圖像分割UNet系列------UNet3+(UNet3plus)詳解

    UNet3+發表于2020年的ICASSP,它是對UNet非常重要的改進,它的性能我認為是可以超過UNet++的,至少在我的使用過程中我會直接使用UNet3+,而不是UNet++。

1、UNet3+結構

    UNet3+主要是參考了UNet和UNet++兩個網絡結構。盡管UNet++采用了嵌套和密集跳過連接配接的網絡結構(見圖1(b)紅色三角區域),但是它沒有直接從多尺度資訊中提取足夠多的資訊。此部分,在我了解而言UNet++雖然名義上通過嵌套和密集跳過連接配接進行了多尺度資訊的利用,但是從本質上看基本都是短連接配接,基本上都對解碼特征進行了再次處理,再加上各個連接配接的融合,多尺度資訊的原始特征幾乎沒有得到特别好的利用,信号處理有些矯枉過正或是丢失。UNet3+利用了全尺度的跳躍連接配接(skip connection)和深度監督(deep supervisions)。全尺度的跳躍連接配接把來自不同尺度特征圖中的進階語義與低級語義直接結合(當然需要必要的上采樣操作);而深度監督則從多尺度聚合的特征圖中學習層次表示。注意一點:UNet++和UNet3+都用到了深度監督,但是監督的位置是完全不一樣的,從圖1(b)、(c)中的Sup部分可以清楚的看到不同之處。

圖像分割UNet系列------UNet3+(UNet3plus)詳解

    在UNet3+中,可以從全尺度捕獲細粒度的細節和粗粒度的語義。為了進一步從全尺寸的聚合特征圖中學習層次表示法,每個邊的輸出都與一個混合損失函數相連接配接,這有助于精确分割,特别是對于在醫學圖像體積中出現不同尺度的器官。從圖1中也可以看出,UNet3+的參數量明顯小于UNet++。

    在介紹Res-UNet和Attention Unet兩個改進算法時我們提到為了滿足醫學圖像分割的準确性要求它們引入注意力機制(Res-UNet為hard-attention,而Attention Unet本質上是soft-attention),當然,在推理時執行預定義的細化方法(如CRF[4])也可以保障圖像分割的準确性。但是,作者并沒有應用這兩種方法,它是利用分類任務來預測輸入圖像是否有器官,為分割任務提供必要的指導。

2、UNet3+主要部分說明

2.1 全尺寸跳躍連接配接(Full-scale Skip Connections)

    全尺寸跳躍連接配接部分的結構如下圖2所示。

圖像分割UNet系列------UNet3+(UNet3plus)詳解

    作者在文章中利用圖2詳細說明了構造 X D e 3 X_{D e}^{3} XDe3​ 特征圖的全過程。從上圖中可以看出,特征圖 X D n 3 X_{D n}^{3} XDn3​的全尺寸連接配接主要是來自三個部分:

  1. 較小尺度編碼器 X E n 1 X_{E n}^{1} XEn1​ 和 X E n 2 X_{E n}^{2} XEn2​。此部分保留了細節資訊,即融合細粒度語義。 X E n 1 X_{E n}^{1} XEn1​ 和 X E n 2 X_{E n}^{2} XEn2​要進行下采樣,分别下采樣4倍和2倍,以統一特征圖的大小。随後進行3*3卷積,64通道。值得注意的一點:此部分的下采樣采用的是無重疊最大池化。
  2. 相同尺度編碼器的特征圖 X E n 3 X_{E n}^{3} XEn3​ 。此層直接進行3*3卷積,64通道操作。
  3. 較大尺度解碼器中的 X D n 4 X_{D n}^{4} XDn4​ 和 X E n 5 X_{E n}^{5} XEn5​ 。此部分融合了粗粒度語義資訊。 X D n 4 X_{D n}^{4} XDn4​與 X E n 5 ( X D e 5 ) X_{E n}^{5}\left(X_{D e}^{5}\right) XEn5​(XDe5​)分别進行2倍與4倍的上采樣,也是為了保證特征圖大小的統一。注意:此部分的上采樣是采用雙線性插值的方法。

    五層疊加(拼接融合)形成320通道的特征圖。随後再進行320通道的3*3卷積、BN、Relu等操作形成新的特征圖 X D n 3 X_{D n}^{3} XDn3​。其他解碼部分的特征圖生成過程類似。

2.2 UNet、UNet++和UNet3+參數數量計算與比較

    在文章開始我們談到UNet3+與UNet++和UNet相比擁有更少的參數,具體怎樣通過計算進行相應的對比呢?論文中,作者進行了相關計算。首先,通過公式表示全尺寸跳躍連接配接, i i i表示沿着編碼的方向第 i i i個下采樣層, N N N表示編碼器個數(論文為5個),那麼特征圖 X D e i X_{D e}^{i} XDei​ 的計算公式如下:

X D e i = { X E n i , H ( [ C ( D ( X E n k ) ) k = 1 i − 1 , C ( X E n i ) ⏟ Scales:  1 t h ∼ i t h , C ( U ( X D e k ) ) k = i + 1 N ⏟ Scales:  ( i + 1 ) t h ∼ N t h ] ) , i = 1 , ⋯   , N − 1 ( 1 ) X_{D e}^{i}=\left\{\begin{array}{l} X_{E n}^{i}, \\ \mathcal{H}\left([\underbrace{C\left(\mathcal{D}\left(X_{E n}^{k}\right)\right)_{k=1}^{i-1}, C\left(X_{E n}^{i}\right)}_{\text {Scales: } 1^{t h} \sim i^{t h}}, \underbrace{C\left(\mathcal{U}\left(X_{D e}^{k}\right)\right)_{k=i+1}^{N}}_{\text {Scales: }(i+1)^{t h} \sim N^{t h}}]\right), i=1, \cdots, N-1 \end{array}\right.(1) XDei​=⎩⎪⎪⎪⎨⎪⎪⎪⎧​XEni​,H⎝⎜⎛​[Scales: 1th∼ith

C(D(XEnk​))k=1i−1​,C(XEni​)​​,Scales: (i+1)th∼Nth

C(U(XDek​))k=i+1N​​​]⎠⎟⎞​,i=1,⋯,N−1​(1)

其中,函數 C ( . ) C(.) C(.)表示卷積操作,函數 H ( . ) \mathcal{H}(.) H(.)表示特征聚合機制 (一個卷積層+一個BN+一個ReLU), 函數 D \mathcal{D} D和 函數 U \mathcal{U} U分别表示上采樣和下采樣操作, [ ] [ ] []表示通道次元拼接融合。

    從圖1可以看出,UNet、UNet++和UNet3+編碼器的結構三者都是一樣的, X E n i X_{E n}^{i} XEni​ 都為 32 × 2 i 32 \times 2^{i} 32×2i 通道數,編碼部分的參數都是一樣多的。他們的不同主要是展現在解碼部分,具體參數量的計算如下所示:

    1)UNet解碼部分

    UNet的解碼部分和編碼部分是對稱的, 是以 X D e i X_{D e}^{i} XDei​ 都為 32 × 2 i 32 \times 2^{i} 32×2i 通道。 UNet解碼部分中第 i i i個解碼階段的參數數量可以通過下面式子計算得到:

P U − D e i = D F × D F × [ d ( X D e i + 1 ) × d ( X D e i ) + d ( X D e i ) 2 + d ( X E n i + X D e i ) × d ( X D e i ) ] ( 2 ) \begin{aligned} P_{U-D e}^{i}=D_{F} \times D_{F} & \times\left[d\left(X_{D e}^{i+1}\right) \times d\left(X_{D e}^{i}\right)+d\left(X_{D e}^{i}\right)^{2}\right.\\ &\left.+d\left(X_{E n}^{i}+X_{D e}^{i}\right) \times d\left(X_{D e}^{i}\right)\right] \end{aligned}(2) PU−Dei​=DF​×DF​​×[d(XDei+1​)×d(XDei​)+d(XDei​)2+d(XEni​+XDei​)×d(XDei​)]​(2)

其中, D F D_{F} DF​ 表示卷積核的大小,函數 d d d表示求出節點的深度。

    2)UNet++解碼部分

    在UNet++中, 它在每一條跳躍路徑 上都利用了稠密卷積子產品(dense conv block),它的計算公式如下:

P U + + − D e i = D F × D F × [ d ( X D e i + 1 ) × d ( X D e i ) + d ( X D e i ) 2 + d ( X E n i + ∑ k = 1 N − 1 − i X M e i , k + X D e i ) × d ( X D e i ) ] ( 3 ) \begin{aligned} P_{U^{++}-D e}^{i}=D_{F} \times D_{F} \times[& d\left(X_{D e}^{i+1}\right) \times d\left(X_{D e}^{i}\right)+d\left(X_{D e}^{i}\right)^{2}+\\ &\left.d\left(X_{E n}^{i}+\sum_{k=1}^{N-1-i} X_{M e}^{i, k}+X_{D e}^{i}\right) \times d\left(X_{D e}^{i}\right)\right] \end{aligned}(3) PU++−Dei​=DF​×DF​×[​d(XDei+1​)×d(XDei​)+d(XDei​)2+d(XEni​+k=1∑N−1−i​XMei,k​+XDei​)×d(XDei​)]​(3)

    從公式看,UNet++的參數量比UNet還要大一些。

    3)UNet3+解碼部分

    在UNet3+中, 每一個解碼器由 N N N個尺度連接配接所成, 是以産生 64 × N 64 \times N 64×N 通道,計算公式如下:

P u 3 + − D e i = D F × D F × [ ( ∑ k = 1 i d ( X E n k ) + ∑ k = i + 1 N d ( X D e k ) ) × 64 + d ( X D e i ) 2 ] ( 4 ) P_{u^{3+}-D e}^{i}=D_{F} \times D_{F} \times\left[\left(\sum_{k=1}^{i} d\left(X_{E_{n}}^{k}\right)+\sum_{k=i+1}^{N} d\left(X_{D e}^{k}\right)\right) \times 64+d\left(X_{D e}^{i}\right)^{2}\right](4) Pu3+−Dei​=DF​×DF​×[(k=1∑i​d(XEn​k​)+k=i+1∑N​d(XDek​))×64+d(XDei​)2](4)

    通過公式(2)(3)(4)可以看出,雖然從網絡結構上看,UNet最為清晰明了,貌似參數應該更少,其實并不是這樣。在保障相同的編碼部分的前提下,它們三者中UNet3+的參數量最少,其次才是UNet,UNet++的參數量是最多的(結構也最複雜)。

2.3 全尺寸深監督(Full-scale Deep Supervision)

    UNet3+全尺寸深監督部分的結構如下圖1所示。當然,UNet3+的深監督應該是從UNet++中得到了啟發。但是,兩者有很大的不同。

    UNet++深監督部分

    UNet++是對第一層的特征圖進行深監督,即對全分辨率特征圖進行深監督,見圖1, X M e 1 , 1 X_{M e}^{1,1} XMe1,1​ 、 X M e 1 , 2 X_{M e}^{1,2} XMe1,2​、 X M e 1 , 3 X_{M e}^{1,3} XMe1,3​、 X D e 1 , 4 X_{D e}^{1,4} XDe1,4​。在實際操作中UNet++使用1*1卷積分别對 X M e 1 , 1 X_{M e}^{1,1} XMe1,1​ 、 X M e 1 , 2 X_{M e}^{1,2} XMe1,2​、 X M e 1 , 3 X_{M e}^{1,3} XMe1,3​、 X D e 1 , 4 X_{D e}^{1,4} XDe1,4​進行操作,去監督每個分支的輸出。

    UNet3+深監督部分

    UNet3+全尺寸深監督是每個解碼器對應一個側輸出(side output),通過ground truth進行監督。為了實作深度監控,每個解碼器的最後一層被送入一個普通的3 × 3卷積層,然後是一個雙線性上采樣和一個sigmoid函數。

    此處進行雙線性上采樣的目的我認為主要有兩個:

  1. 上采樣是将第2、3、4、5層擴充成全分辨率特征圖,保證與第一層相同,這也是全尺寸深監督的關鍵操作;
  2. 雙線性上采樣的方式可以最大限度保證上采樣過程中邊緣資訊的完整性(醫學圖像邊緣的不确定性決定要盡量保障邊緣資訊不丢失);

    UNet3+的Loss函數

    為了進一步增強器官的邊界,作者提出了一個基于多尺度結構相似度指數(MS-SSIM)的損失函數來賦予模糊邊界更高的權重。受益于此,UNet 3+将關注模糊邊界,區域分布差異越大,MS-SSIM值越高。從分割結果 P P P和 ground truth mask G G G裁剪了兩個相應的 N × N N×N N×N大小的塊,它們可以分别表示為

= { p j : j = 1 , … , N 2 } = 1 , … , N 2 } \left.=\left\{p_{j}: j=1, \ldots, N^{2}\right\}=1, \ldots, N^{2}\right\} ={pj​:j=1,…,N2}=1,…,N2} 和 g = { g j : j = 1 , … , N 2 } g=\left\{g_{j}: j=1, \ldots, N^{2}\right\} g={gj​:j=1,…,N2} 。定義 p p p和 g g g 的MS-SSIM損失函數為:

ℓ m s − s s i m = 1 − ∏ m = 1 M ( 2 μ p μ g + C 1 μ p 2 + μ g 2 + C 1 ) β m ( 2 σ p g + C 2 σ p 2 + σ g 2 + C 2 ) γ m \ell_{m s-s s i m}=1-\prod_{m=1}^{M}\left(\frac{2 \mu_{p} \mu_{g}+C_{1}}{\mu_{p}^{2}+\mu_{g}^{2}+C_{1}}\right)^{\beta_{m}}\left(\frac{2 \sigma_{p g}+C_{2}}{\sigma_{p}^{2}+\sigma_{g}^{2}+C_{2}}\right)^{\gamma_{m}} ℓms−ssim​=1−m=1∏M​(μp2​+μg2​+C1​2μp​μg​+C1​​)βm​(σp2​+σg2​+C2​2σpg​+C2​​)γm​

其中, M M M表示尺度的總數量, μ p , μ g \mu_{p}, \mu_{g} μp​,μg​ 和 σ p , σ g \sigma_{p}, \sigma_{g} σp​,σg​ 分别是 p p p、 g g g的均值和方差, σ p g \sigma_{p g} σpg​ 表示它們的協方差。 β m , γ m \beta_{m}, \gamma_{m} βm​,γm​ 定義每個尺度中這兩個部分的相對重要性,設定可以參考[9]。加入兩個小的常量 C 1 = 0.0 1 2 C_{1}=0.01^{2} C1​=0.012 和 C 2 = 0.0 3 2 C_{2}=0.03^{2} C2​=0.032 是避免除零的不穩定情況。作者将尺度設定為5(和UNet、UNet++保持一緻)。

    結合focal損失函數 ℓ f l \ell_{f l} ℓfl​[10], MS-SSIM損失函數 ℓ m s − s s i m \ell_{m s-s s i m} ℓms−ssim​和IoU損失函數 ℓ i o u \ell_{i o u} ℓiou​[11],我們提出一種混合損失函數用于像素級别分割、塊分割、圖像級别分割三個不同層次的分割,它可以同時擷取大尺度和精細結構中清晰的邊緣。混合分割損失被定義為:

ℓ s e g = ℓ f l + ℓ m s − s s i m + ℓ i o u \ell_{s e g}=\ell_{f l}+\ell_{m s-s s i m}+\ell_{i o u} ℓseg​=ℓfl​+ℓms−ssim​+ℓiou​

2.4 Classification-guided Module (CGM)

    在這個子產品中,主要是利用分類任務來解決由于較淺層中背景噪聲等殘留資訊導緻過分割的現象。CGM的具體結果如下圖所示:

圖像分割UNet系列------UNet3+(UNet3plus)詳解

    如圖3所示,最深層次的二維張量 X E n 5 X_{E n}^{5} XEn5​ 經過一系列的操作包括dropout、1*1卷積、maxpooling、sigmoid,最後有兩個值代表有/沒有器官的機率。利用最豐富的語義資訊,分類結果可以進一步指導每一個切分側邊輸出兩個步驟。首先,在Argmax函數的幫助下,将二維張量轉化為 { 0 , 1 } \{0,1\} {0,1} 的單個輸出,表示有/沒有器官。随後,我們将單個分類輸出與側分割輸出相乘。由于二分類任務的簡單性,該子產品通過優化交叉熵損失函數,輕松獲得準确的分類結果,避免對非器官圖像過分割。

    從圖3可以看出,分類部分主要是利用第五層,第五層是尺度最大的層,分類結果相對與其它層肯定會粗糙一些,因為用第五層判斷第一層是不是器官,必然會有上采樣或是第一層下采樣的過程。但是,同樣因為這個原因,第五層分類不會受到太多細節資訊的影響,反而會使分類的區域定位也更加準确,這是一個非常巧妙的設計。

3、總結

    UNet3+的優點和主要貢獻在文章中都穿插講述了,在總結部分主要提一下我自在實際應用中遇到的一些問題。

  1. 輸入圖像大小問題。UNet3+在編碼部分采用的是無重疊的最大池化,這就要求圖像下采樣過程中的尺寸必須是偶數,在實際操作中可以先考慮最後一層(第5層),例如從下倒上反推依次為8、16、32、64、128。
  2. loss應用問題。作者使用了focal損失函數 ℓ f l \ell_{f l} ℓfl​, MS-SSIM損失函數 ℓ m s − s s i m \ell_{m s-s s i m} ℓms−ssim​和IoU損失函數 ℓ i o u \ell_{i o u} ℓiou​的混合損失函數,但是我在實際使用中并沒有嚴格遵守論文的loss,而是對五層輸出的loss進行了權重操作并以第一層的loss為主。當然,總體效果還是不錯的。
  3. 資料問題。由于自己項目中的圖像品質總體偏差,是以原始資料進行了圖像增強(數字圖像處理定義中的增強,而不是增加樣本)。從處理前後的資料訓練結果對比來看,提高了3個百分點。

參考:https://zhuanlan.zhihu.com/p/136164721

繼續閱讀