概述

在這篇論文中，提出了一種新的醫(yī)學(xué)圖像分割混合架構(gòu)：PHTrans，它在主要構(gòu)建塊中并行混合 Transformer 和 CNN，分別從全局和局部特征中生成層次表示并自適應(yīng)聚合它們，旨在充分利用 Transformer 和 CNN 各自的優(yōu)勢以獲得更好的分割性能。

具體來說，PHTrans 沿用 U 形設(shè)計，在深層引入并行混合模塊，其中卷積塊和修改后的 3D Swin Transformer 塊分別學(xué)習(xí)局部特征和全局依賴關(guān)系，然后使用 sequence-to-volume 操作統(tǒng)一輸出維度以實(shí)現(xiàn)特征聚合，操作的具體細(xì)節(jié)在這篇閱讀筆記的后面詳細(xì)介紹。最后在 BCV 和 ACDC 數(shù)據(jù)集上驗(yàn)證了其有效性，并用 nnUNet 包預(yù)處理 BCV 和 ACDC 數(shù)據(jù)集。

為什么要并行

下圖的 (a)~(d) 是幾種流行的基于 Transformer 和 CNN 的混合架構(gòu)，既將 Transformer 添加到以 CNN 為 backbone 的模型中，或替換部分組件。其中 (c) 與 (b) 的區(qū)別是通過 Transformer 橋接從編碼器到解碼器的所有階段，而不僅僅是相鄰的階段，這就捕獲了多尺度全局依賴。(d) 表示將 Transformer 和 CNN 交織成一個混合模型，其中卷積編碼精確的空間信息，而自注意力機(jī)制捕獲全局上下文信息。

圖 (e) 表示二者的并行。在串行組合中，卷積和自注意力機(jī)制無法貫穿整個網(wǎng)絡(luò)架構(gòu)，難以連續(xù)建模局部和全局表示，因此這篇論文里認(rèn)為并行可以充分發(fā)揮它們的潛力。

PHTrans 架構(gòu) overview

首先，我們從總體上分析一下 PHTrans 架構(gòu)，然后在下一部分看它的細(xì)節(jié)。如下圖 (b)，其主要構(gòu)建塊由 CNN 和 Swin Transformer 組成，以同時聚合全局和局部表示。圖 (a) 依舊遵循的 U 形架構(gòu)設(shè)計，在淺層只是普通的卷積塊，在深層引入了 sequence-to-volume 操作來實(shí)現(xiàn) Swin Transformer 和 CNN 在一個塊中的并行組合。

我們上一篇解析的 UNeXT 也是只在深層使用 TokMLP 的，看來淺層的卷積還是必要的。也就是說，與串行混合架構(gòu)相比，PHTrans 可以獨(dú)立并行構(gòu)建分層的局部和全局表示，并在每個階段融合它們。

進(jìn)一步解釋下為什么輸入的第一層也就是 U 型架構(gòu)的淺層沒有用 Trans&Conv Block？因?yàn)樽宰⒁饬C(jī)制的計算復(fù)雜度高，Transformer 無法直接接收以像素為標(biāo)記的輸入。在論文的實(shí)現(xiàn)中，使用了級聯(lián)卷積塊和下采樣操作來減小空間大小，逐步提取高分辨率的低級特征以獲得精細(xì)的空間信息。類似地，這些純卷積模塊也部署在解碼器的對應(yīng)層，并通過上采樣恢復(fù)原始維度。

我們仔細(xì)看下 PHTrans 的編碼器，對于 H×W×D 的輸入 volume（3D 醫(yī)學(xué)圖像），其中 H、W 和 D 分別表示高度、寬度和深度，首先使用幾個純卷積模塊得到的 volume，其中 N1 和 C 表示卷積塊和通道的數(shù)量。然后輸入到 Trans&Conv Block 重復(fù) N2 次。對于解碼器同樣基于純卷積模塊和并行混合模塊構(gòu)建，并通過跳躍連接和加法操作融合來自編碼器的語義信息。此外，在訓(xùn)練期間在解碼器的每個階段都使用深度監(jiān)督機(jī)制，產(chǎn)生總共 N1 + N2 個輸出，其中應(yīng)用了由交叉熵和 DICE 的聯(lián)合損失。

深度監(jiān)督（deep supervision）又稱為中繼監(jiān)督（intermediate supervision），其實(shí)就是網(wǎng)絡(luò)的中間部分新添加了額外的 Loss，跟多任務(wù)是有區(qū)別的，多任務(wù)有不同的 GT 計算不同的 Loss，而深度監(jiān)督的 GT 都是同一個 GT，不同位置的 Loss 按系數(shù)求和。深度監(jiān)督的目的是為了淺層能夠得到更加充分的訓(xùn)練，避免梯度消失（有待研究）。在提供的 Github 代碼里，提到的超參數(shù)有 N1、N2、M1 和 M2，M1 和M2 是并行混合模塊中 Swin Transformer 塊和卷積塊的數(shù)量。

Trans&Conv block

Trans&Conv block 的設(shè)計是我們最感興趣的地方。縮小比例的特征圖分別輸入 Swin Transformer (ST) 塊和卷積 (Conv) 塊，分別在 ST 塊的開頭和結(jié)尾引入 Volume-to-Sequence (V2S) 和 Sequence-to-Volume (S2V) 操作來實(shí)現(xiàn) volume 和 sequence 的變換，使其與 Conv 塊產(chǎn)生的輸出兼容。具體來說，V2S 用于將整個 3D 圖像重塑為具有窗口大小的 3D patches 序列。S2V 是相反的操作。

如上一節(jié)的圖 (b) 所示，一個 ST 塊由一個基于移位窗口的多頭自注意力 (MSA) 模塊組成，然后是一個 2 層 MLP。在每個 MSA 模塊和每個 MLP 之前應(yīng)用一個 LayerNorm (LN) 層，在每個模塊之后應(yīng)用一個殘差連接。在 M1 個連續(xù)的 ST 塊中，W-MSA 和 SW-MSA 交替嵌入到 ST 塊中，W-MSA能夠降低計算復(fù)雜度，但是不重合的窗口之間缺乏信息交流，這樣其實(shí)就失去了 Transformer 利用 Self-Attention 從全局構(gòu)建關(guān)系的能力，于是用 SW-MSA 來跨窗口進(jìn)行信息交流（跨窗口連接），同時保持非重疊窗口的高效計算。

對于醫(yī)學(xué)圖像分割，需要將標(biāo)準(zhǔn) ST 塊修改為 3D 版本，該版本在局部 3D 窗口內(nèi)計算自注意力，這些窗口被安排為以非重疊方式均勻劃分體積。計算方法是下面這樣的：假設(shè) x ∈ H×W×S×C 是 ST 塊的輸入，首先將其 reshape 為 N×L×C，其中 N 和 L = Wh × Ww × Ws 分別表示 3D 窗口的數(shù)量和維度。每個 head 中的 self-attention 計算如下：

Q, K, V ∈ L×d 是查詢、鍵和值矩陣，d 是查詢/鍵維度，B ∈ L×L 是相對位置偏差。B 的取值在論文和代碼里都可以找到，這里我們就不仔細(xì)探究了。(b) 中的卷積塊以 3 × 3 × 3 卷積層、GELU 非線性和實(shí)例歸一化層為單位重復(fù) M2 次。最后，通過加法運(yùn)算融合 ST 塊和 Conv 塊的輸出。編碼器中 Trans&Conv 塊的計算過程（抽象成并行）可以總結(jié)如下：

xi?1 是編碼器第 i?1 階段的下采樣結(jié)果。值得注意的是，在解碼器中，除了跳躍連接之外，還通過加法操作來補(bǔ)充來自編碼器的上下文信息（圖 (a) 中的圈 C 和 圈 +）。因此，解碼器中的 Trans&Conv 塊計算（抽象成并行）可以表示為：

實(shí)驗(yàn)

實(shí)驗(yàn)在 BCV 和 ACDC 數(shù)據(jù)集上，BCV 分割腹部 CT 多個目標(biāo)，ACDC 是 MRI 心臟分割，標(biāo)記了左心室 (LV)、右心室 (RV) 和心肌 (MYO)。在 BCV 上和其他 SOTA 方法的比較如下表：

在 ACDC 上和其他 SOTA 方法的比較如 Table 2 所示，Table 3 中的參數(shù)量和 FLOPS 和其他方法比也沒有很夸張，參數(shù)量甚至和 nnU-Net 相近。

可視化分割結(jié)果如下圖，我們只定位藍(lán)色肝臟的分割效果，箭頭位置表明分割的效果 PHTrans 是更優(yōu)秀的。

總結(jié)

PHTrans 也許為更多下游醫(yī)學(xué)圖像任務(wù)開發(fā)了新的可能性。在 PHTrans 中，都是普通的 Swin Transformer 和簡單的 CNN 塊，這表明性能提升源于并行混合架構(gòu)設(shè)計，而不是 Transformer 和 CNN 塊。此外，PHTrans 沒有經(jīng)過預(yù)訓(xùn)練，因?yàn)榈侥壳盀橹惯€沒有足夠大的通用 3D 醫(yī)學(xué)圖像數(shù)據(jù)集。