楊 鶴，柏正堯

云南大學 信息學院，昆明 650500

目前，基于CNN的方法在醫(yī)學圖像分割領(lǐng)域取得了優(yōu)異的成績?，F(xiàn)有的醫(yī)學圖像分割方法主要依賴于U形結(jié)構(gòu)的全卷積神經(jīng)網(wǎng)絡[1-3]。其中最典型的U形結(jié)構(gòu)網(wǎng)絡，U-Net[1]由一個帶跳躍連接的對稱編碼器-解碼器組成。該網(wǎng)絡為了減少下采樣過程中造成的空間信息丟失。通過跳躍連接來增強細節(jié)保留。并在許多醫(yī)學圖像分割任務中獲得了巨大的成功，占據(jù)著主導地位。按照這一思想，出現(xiàn)了許多U-Net的變體，如3D UNet[4]、Res-UNet[5]、U-Net++[6]、U-Net3+[7]用于各種醫(yī)學圖像分割。

在上述網(wǎng)絡中，由于卷積運算的固有局限性，純CNN的方法很難學習明確的全局和遠程語義信息交互[8]。為了克服這樣的局限性，現(xiàn)有研究建議建立基于CNN特征的自注意機制[9-10]。Transformer在建模全局上下文方面功能強大[8]。但是純Transformer的方法也存在特征丟失的問題。于是，基于CNN架構(gòu)與Transformer模塊相結(jié)合的方法得到了越來越廣泛的關(guān)注。例如，文獻[11]利用自注意機制來增強卷積算子，將強調(diào)局部性的卷積特征映射與能夠?qū)θ诌M行建模的自注意特征映射連接起來。DETR[12]使用一個傳統(tǒng)的CNN骨架網(wǎng)絡來學習輸入圖像的二維表示，編碼器和解碼器都是由Transformer構(gòu)成。其中，TransUNet[8]則將這種方法應用到醫(yī)學圖像分割領(lǐng)域，同時具有Transformer和CNN的優(yōu)點。將Transformer作為醫(yī)學圖像分割任務的強編碼器，與CNN相結(jié)合，通過恢復局部空間信息來增強更精細的細節(jié)[8]。但TransUNet采用傳統(tǒng)CNN作為特征提取器和上采樣器，在特征提取和上采樣階段，卷積中感受野太小的局限性依然存在。其中，Swin-UNet[13]，一種基于純Transformer的U-Net形的醫(yī)學圖像分割網(wǎng)絡。把標記化的圖像塊通過跳躍連接饋送到基于Transformer的U形En-Decoder架構(gòu)中，以進行局部和全局語義特征學習。但是純Transformer的方法，存在特征丟失的問題。ViT[14]可以像Transformer一樣直接通過自注意機制處理圖像塊或CNN輸出。此外，文獻[15]還提出了一種局部自注意模塊，可以完全取代ResNet架構(gòu)中的3×3卷積。但是這些設計都忽略了相鄰keys之間的豐富上下文。于是出現(xiàn)了一種新型的Transformer式模塊，CoTNet[16]，其充分利用了輸入keys之間的上下文信息來指導動態(tài)注意矩陣的學習[16]。在傳統(tǒng)自注意學習基礎(chǔ)上，將局部上下文信息與全局上下文信息融合起來，更好地增強視覺表達能力。

為了解決現(xiàn)有醫(yī)學圖像分割網(wǎng)絡中卷積的感受野太小以及對局部上下文信息利用不足的問題。提升現(xiàn)有醫(yī)學圖像分割網(wǎng)絡的精度和速度。CoT-TransUNet在編碼器部分，采用CoTNet作為特征提取器以提取更豐富特征。并把Transformer的層數(shù)增加到了16層。Transformer塊把特征圖編碼為輸入序列后送入解碼器中進行上采樣。在解碼器部分，上采樣器中采用具有更大感受野的CARAFE算子，在保持輕量化的同時又能基于內(nèi)容上采樣。最后，通過跳躍連接實現(xiàn)編碼器和解碼器在不同分辨率上的特征聚合。在TransUNet的基礎(chǔ)上結(jié)合以上改進，更好地提升了網(wǎng)絡的性能。在多器官分割任務中，以DSC和HD為評價指標，對8個腹部器官（主動脈、膽囊、脾臟、左腎、右腎、肝、胰腺、脾臟、胃）進行評價，并在可視化結(jié)果中用不同的顏色對8個腹部器官進行標注。CoT-TransUNet相對于TransUNet，DSC提升了約0.9%，HD提升了約5.7%。

1 相關(guān)工作

1.1 醫(yī)學圖像分割中CNN與Transformer的結(jié)合

為了解決卷積運算感受野太小而對全局建模產(chǎn)生的不利影響[8]。Chen等人提出了一種新網(wǎng)絡，TransUNet[8]用于代替原有醫(yī)學圖像分割方法中的U-Net[1]和U-Net的一些變體。TransUNet從序列到序列預測的角度建立了自注意機制[8]。但是Transformer也存在特征分辨率損失的問題。于是，TransUNet采用CNN和Transformer相結(jié)合的思想，利用來自CNN特征的詳細高分辨率空間信息和Transformer編碼的全局上下文[8]，并采用U形設計。先使用CNN進行特征提取，然后將提取到的特征映射中的標記化圖像塊編碼為用于提取全局上下文的輸入序列[8]。再對編碼后的特征上采樣并與之前高分辨率的CNN特征結(jié)合。既在一定程度上解決了卷積的局限性，又防止了特征的丟失。在多器官分割任務中，性能也優(yōu)于之前的架構(gòu)。

1.2 基于CNN架構(gòu)的Transformer風格模塊

為了解決Transformer特征丟失的問題，以及更充分地利用局部上下文信息。Li等人提出了一種新的Transformer風格的模塊，即上下文Transformer（CoT）塊[16]。如圖1所示。首先通過k×k卷積對輸入K進行上下文編碼，從而實現(xiàn)輸入的局部（靜態(tài)）上下文表示。然后將K與Q合并在一起，再通過兩個連續(xù)的1×1卷積學習動態(tài)多頭注意矩陣[16]。此矩陣是由Q的信息和局部上下文信息交互得到的，而不只是建模了Q和K之間的關(guān)系。也就是說通過局部上下文信息的引導，增強自注意機制。然后再與特征映射后的V相乘，得到輸入的全局（動態(tài)）上下文表示。最后將靜態(tài)和動態(tài)上下文表示進行線性融合得到輸出。使用CoT塊替換ResNet結(jié)構(gòu)中的3×3卷積后得到CoTNet，在圖像識別，目標檢測和實例分割實驗中，性能都得到了提升。

圖1 CoT塊結(jié)構(gòu)圖Fig.1 CoT block structure diagram

1.3 CARAFE上采樣算子

針對目前應用最廣泛的最近鄰插值法和雙線性插值法，以及反卷積的上采樣方法。前兩個主要關(guān)注亞像素鄰域，無法捕捉豐富語義信息。而反卷積的方法忽略了低層的內(nèi)容，無法對局部變化有很好的反應能力，并且需要大量參數(shù)。Wang等人提出了一種CARAFE[17]上采樣算子。其通過加權(quán)組合在以每個位置為中心的預定義區(qū)域內(nèi)重組特征，以內(nèi)容感知的方式生成權(quán)重，每個位置都有多組這樣的上采樣權(quán)重。然而，這些權(quán)重并不作為網(wǎng)絡參數(shù)學習，而是使用一個輕量級的，帶有softmax激活函數(shù)的全卷積模塊來實時預測。因此，能在大接收域內(nèi)聚合信息的同時，又能實時適應特定實例的內(nèi)容，并能保持計算效率。在不同的任務中，CARAFE上采樣算子僅通過一小部分額外參數(shù)和計算工作就取得了顯著的進步。在CoT-TransUNet中，每個CARAFE上采樣塊都由一個CARAFE上采樣算子、一個3×3卷積、一個ReLU層組成，并級聯(lián)多個這樣的上采樣塊構(gòu)成解碼器。

2 方法

CoT-TransUNet的總體結(jié)構(gòu)如圖2所示。對于一個輸入圖像，首先將其送入混合編碼器模塊并由CoTNet提取特征，接著通過Transformer塊將特征圖編碼成用于提取全局上下文的輸入序列，然后采用上采樣模塊對編碼后的特征上采樣并通過跳躍連接與編碼器中的高分辨率特征結(jié)合，恢復到原始分辨率，以設計一種輕量化的端到端U形網(wǎng)絡結(jié)構(gòu)，從而充分利用特征提取階段的局部上下文信息，增加上采樣階段的感受野以及增強細節(jié)保留。通過這三個模塊的結(jié)合，并從分割結(jié)果可以看出，CoT-TransUNet實現(xiàn)了更好的邊緣預測，獲得了優(yōu)于其他網(wǎng)絡的分割性能。

圖2 CoT-TransUNet的總體框架Fig.2 Overall framework of CoT-TransUNet

2.1 混合編碼器

2.1.1 特征提取

CoTNet：如圖1所示，對于輸入的特征圖，keys、queries、values分別定義為K=X、Q=X和V=XWv。其中，Wv是V的權(quán)值矩陣。在k×k空間網(wǎng)格內(nèi)對所有相鄰keys進行k×k組卷積，所得到的K1∈RH×W×C代表了相鄰keys之間的局部上下文信息，于是將K1作為輸入X的靜態(tài)上下文表示。然后將和Q合并，再通過兩個連續(xù)的1×1卷積（帶ReLU激活函數(shù)的和不帶激活函數(shù)的）獲得注意矩陣D[16]：

然后，把獲得的注意力矩陣D與V點乘，得到，代表了全部上下文信息，于是將作為輸入的動態(tài)上下文表示：

最后，將靜態(tài)上下文表示和動態(tài)上下文表示相加，得到輸出。

用CoTNet-50作為特征提取器，生成特征圖后輸入到Transformer塊中編碼成序列。

2.1.2 圖像序列化

Transformer塊：對于一張圖像X∈ΦH×W×C，其空間分辨率為H×W，通道數(shù)為C。將輸入X重構(gòu)為一些2D的，其中每個patch的大小為P×P，，即輸入序列的長度。將已經(jīng)矢量化的XP映射到隱藏的N維空間中[8]。再把特定的位置嵌入添加到patch嵌入中以保留位置信息，patch嵌入表示如下面公式所示：

其中，R∈Φ(P2·C)×N是patch的嵌入投影，Rpos∈ΦM×N是位置嵌入。

Transformer層是由L層的多頭自注意（MSA）和多層感知器（MLP）塊組成。因此，第λ層的輸出表示為：

其中，LN(·)為層歸一化算子，yL為編碼后的圖像表示。如圖2（1）所示。

通過上述方法便可以得到patch的嵌入。CoTTransUNet使用CoTNet為輸入生成特征映射。再對從CNN特征圖中提取的1×1個patch進行嵌入，而不是對原始圖像進行嵌入[8]。

2.2 解碼器

解碼器部分是由一個級聯(lián)上采樣器構(gòu)成的，將從編碼器輸出的隱藏特征的序列重構(gòu)為后，通過級聯(lián)多個CARAFE上采樣塊將分辨率恢復到H×W的全分辨率。如圖2（2）所示。

上采樣核預測模塊：對于形狀為H×W×C的輸入特征圖，為了減少后續(xù)步驟的參數(shù)和計算量，首先用一個1×1的卷積將通道數(shù)壓縮到Cm。然后把第一步中壓縮后的特征圖，利用一個kencoder×kencoder的卷積層來預測上采樣核，上采樣核尺寸為kup×kup，輸出的通道數(shù)為Cup=σ2ku2p。在此過程中，該編碼器的參數(shù)為kencoder×kencoder×Cm×Cup，可以看出，增加kencoder可以擴大編碼器的感受野，并在更大的區(qū)域內(nèi)利用上下文信息，但是，計算復雜度也會隨著核大小的平方而增加，而更大的核帶來的好處卻不會增加。因此，kencoder=kup-2是性能和效率之間的一個很好的權(quán)衡[17]。然后把通道維在空間維展開，得到形狀為σH×σW×ku2p的上采樣核。并利用softmax進行歸一化，使得卷積核權(quán)重和為1。

特征重組模塊：對于輸出特征圖中的每個位置，將其映射回輸入特征圖，取出以該位置為中心的kup×kup的區(qū)域，并與預測出的該點上采樣核作點積，得到輸出值。達到相同位置的不同通道共享同一個上采樣核的目的。并且，由于內(nèi)核歸一化，CARAFE沒有進行任何縮放和改變特征圖的均值。因為來自局部區(qū)域相關(guān)點的信息得到了更多的關(guān)注，重組后的特征圖擁有更強的語義信息。

2.3 評價指標

對于分割問題，采用Dice相似系數(shù)（Dice similariy coefficient，DSC）和Hausdorff距離來評估模型的分割性能。并通過對參數(shù)數(shù)量、計算復雜度（GFLOPs）和推理時間（ms）的比較來評估模型的輕量化。

2.3.1 Dice系數(shù)

Dice系數(shù)是衡量兩個集合相似度的指標，用于計算兩個樣本的相似度，取值范圍是[0，1]，值越大說明分割結(jié)果與Ground Truth越接近，分割效果越好。如下所示：

式中，A和B分別表示Ground Truth和預測的結(jié)果。

2.3.2 Hausdorff距離

Hausdorff距離是在度量空間中任意兩個集合之間定義的一種距離。表示分割結(jié)果與Ground Truth兩個點集之間最短距離的最大值。如下所示：

式中，d(a,b)表示a、b兩點之間的歐氏距離。

3 實驗

3.1 數(shù)據(jù)集

使用的數(shù)據(jù)集是MICCAI 2015多圖譜腹部標記挑戰(zhàn)賽中的Synapse多器官CT數(shù)據(jù)集。此數(shù)據(jù)集收集了30例患者，一共3 779張軸向腹部臨床CT圖像。按照文獻[8，18]，隨機將18個樣本（2 211張軸向切片）做為訓練集，12個樣本做為測試集。測量其DSC（平均Dice相似系數(shù)，單位：%）和HD（平均Hausdorff距離，單位：mm）作為最終的評價指標。并采用DSC作為評價指標，分別對8個腹部器官（主動脈、膽囊、左腎、右腎、肝臟、胰腺、脾臟、胃）進行評價。

3.2 實施細節(jié)

整個實驗是在Python 3.7和Pytorch 1.8.0上實現(xiàn)的。實驗開始階段都對數(shù)據(jù)進行了隨機翻轉(zhuǎn)和旋轉(zhuǎn)的預處理，以增強數(shù)據(jù)的多樣性。輸入圖像的大小設置為224×224，patch size設置為16，batch size設置為24，在解碼器部分連續(xù)級聯(lián)四個CARAFE上采樣塊以達到完全分辨率?？蚣懿捎肅oT-TransUNet-50。在編碼器設計中，將CoTNet-50、CoTNet-101與ViT[14]組合，記作C50-ViT和C101-ViT，并且都在ImageNet[19]進行了預訓練。對于模型，使用交叉熵損失（cross entropy loss，CE）和Dice損失（Dice loss）的組合來訓練CoT-TransUNet-50，如下式所示：

并且采用學習率為0.01，動量為0.9，權(quán)重衰減為1E-4的SGD優(yōu)化器進行訓練，迭代的默認次數(shù)為14 000次。實驗是在單個12 GB內(nèi)存的Nvidia RTX3060 GPU上進行的。

3.3 實驗結(jié)果

在Synapse多器官分割數(shù)據(jù)集上進行實驗，并與6個之前的先進技術(shù)：U-Net[1]、Swin-UNet[13]、V-Net[20]、Att-UNet[21]、DARR[18]、TransUNet[8]進行了對比結(jié)果如表1所示。CoT-TransUNet-50獲得了最好的分割性能，分割精度達到了78.24%（DSC）和23.75 mm（HD）。與TransUNet相比DSC評價指標提高了約0.9個百分點，雖然改進不大，但是在HD評價指標上卻提高了約5.7 mm。如表2所示，在與目前幾種主流的網(wǎng)絡進行比較后，相比于其他引入注意力機制的網(wǎng)絡，例如Swin-UNet[13]、TransUNet[8]。CoT-TransUNet-50的參數(shù)數(shù)量更少，推理時間更快，計算復雜度也更低，也更加輕量化。不同框架在Synapse多器官CT數(shù)據(jù)集的分割結(jié)果如圖3所示。從圖中可以看出，基于純CNN的方法對邊界信息不夠敏感，容易出現(xiàn)過度分割的問題。例如，在第一行和第二行中，胃部被TransUNet[8]、Att-UNet[21]、U-Net[1]過度分割?；诩僒ransformer的方法雖然對邊界信息較為敏感（特別是胃部的分割），但是由于一些特征的丟失，容易出現(xiàn)分割不足的問題。例如，在第二行中，Swin-UNet[13]對胰腺的預測出現(xiàn)缺失，而CoT-TransUNet則正確預測了胰腺，并保留了很好的邊界信息。實驗結(jié)果表明，相比于TransUNet，以及其他基于純CNN的框架，CoT-TransUNet更加注重邊界信息，能實現(xiàn)更好的邊緣預測。對于純Transformer的方法，CoT-TransUNet既保證了對邊界信息的敏感度，又防止了特征的丟失。在傳統(tǒng)CNN中加入自注意機制[22-23]和Transformer后，將局部上下文信息和全局上下文信息相融合，提高了輸出特征的表達能力，再通過Transformer編碼以及U形結(jié)構(gòu)的上采樣恢復，從而獲得更好的分割結(jié)果。

表1 Synapse多器官CT數(shù)據(jù)集上不同網(wǎng)絡的分割精度對比Table 1 Comparison of segmentation accuracy under different networks on Synapse multi-organ CT dataset

表2 與幾種主流網(wǎng)絡的對比Table 2 Compares with several major networks

圖3 不同框架的分割結(jié)果Fig.3 Segmentation results of different frames

3.4 分析研究

為了全面評估CoT-TransUNet框架的性能，并探究不同因素對其性能的影響。于是在Synapse數(shù)據(jù)集上對CoT-TransUNet進行了各種研究。包括結(jié)構(gòu)的設計、輸入圖像的分辨率、CNN的深度、patch size的大小、模型的大小。

結(jié)構(gòu)的設計：為了分析CoT塊在編碼器部分的重要性，通過在編碼器、解碼器、編-解碼器中分別加入CoT塊，以及不加入CoT塊進行實驗。實驗結(jié)果如表3所示。在編碼器中加入CoT塊，即原網(wǎng)絡中的CoTNet-Transformer混合編碼器結(jié)構(gòu)，可以獲得更好的分割精度。為了分析CARAFE上采樣算子在解碼器部分的重要性，在解碼器中分別采樣CARAFE算子和2×上采樣算子進行實驗。結(jié)果如表4所示。在解碼器中采用CARAFE上采樣算子相比于2×上采樣算子可以獲得更好的分割精度。

表3 CoT塊替換位置的消融研究Table 3 Ablation study of replacement position of CoT block 單位：%

表4 上采樣算子的消融研究Table 4 Ablation study of upsampling operator單位：%

輸入圖像的分辨率：為了探究更高分辨率對分割性能的影響。在其他參數(shù)不變的情況下，對分辨率為512×512的圖像進行實驗。實驗結(jié)果如表5所示。在保持patch size不變的情況下，增大分辨率導致Transformer的序列長度增加。文獻[24-25]指出，增加有效序列長度體現(xiàn)出對魯棒的改進。雖然模型的分割精度得到提升，但是計算成本過大。所以所有實驗都是在默認分辨率（224×224）下進行。

表5 輸入圖像分辨率的研究Table 5 Study of input image resolution 單位：%

CNN的深度：由于卷積神經(jīng)網(wǎng)絡的靈活性，為了探究是否能推廣到更深的網(wǎng)絡層次。于是對CoT-TransUNet-101進行實驗并與CoT-TransUNet-50比較。實驗結(jié)果如表6所示。CoT-TransUNet-101并沒有獲得優(yōu)于CoTTransUNet-50的分割精度。

表6 CNN深度的研究Table 6 Study of CNN depth 單位：%

patch size的大小：因為Transformer的序列長度與patch size的平方成反比[26]。patch size越小，序列長度越大。為了探究patch size對分割性能的影響。分別對24、16、8的patch size進行實驗。實驗結(jié)果如表7所示。減小patch size確實可以起到提升分割性能的作用[27-28]。按照TransUNet[8]、ViT[14]中的設置，CoT-TransUNet使用16×16作為默認值。

表7 Patch size大小的研究Table 7 Study of patch size 單位：%

模型的規(guī)模：最后，對不同模型大小的CoT-TransUNet進行了消融研究。主要有兩種不同的CoT-TransUNet配置?！盎A(chǔ)版”和“升級版”，按照TransUNet[8]、ViT[14]，對于“基礎(chǔ)板”，hidden size為768，Transformer層數(shù)為16，MLP size為3 072，注意力頭數(shù)為12?！吧壈妗敝校琱idden size為1 024，Transformer層數(shù)為24，MLP size為4 096，注意力頭數(shù)為16。如表8所示。由結(jié)果可知，模型越大，效果越好。但是計算成本的增加過大。所以，所有實驗都采用“基礎(chǔ)版的模型”。

表8 模型大小的研究Table 8 Study of model size 單位：%

3.5 討論

在有Transformer參與的模塊中，其性能受到預訓練模型的嚴重影響。由于Google和CoTNet[16]都沒有提供“CoTNet50/101-ViT”的預訓練權(quán)重，只有采用退而求其次的方法，把CoTNet[16]在ImageNet上訓練后提供的單個CoTNet-50/101的預訓練權(quán)重和Google提供的ViT[14]預訓練權(quán)重結(jié)合。由于這兩個模塊的權(quán)重不是從端到端的訓練生成的，一定程度上影響了CoTTransUNet的性能。其次，實驗中所用到的圖像都是2D的，然而，醫(yī)學圖像的數(shù)據(jù)大多數(shù)是3D的。在未來的研究中繼續(xù)探索，將CoT-TransUNet推廣到3D醫(yī)學圖像分割中。

4 結(jié)論

針對現(xiàn)有醫(yī)學圖像分割網(wǎng)絡中卷積運算感受野太小以及對局部上下文信息利用不足，所提出的CoTTransUNet在多器官分割任務中性能優(yōu)于現(xiàn)有網(wǎng)絡。在CoT-TransUNet中，首先，采用更注重局部上下文信息的CoTNet作為特征提取主干，為輸入生成更好的特征圖。其次，采用感受野更大并基于內(nèi)容上采樣的輕量化CARAFE算子進行上采樣。最后，通過跳躍連接將編碼器和解碼器相連，以增強細節(jié)保留。在多器官分割任務的實驗中也表現(xiàn)出更好的分割精度和更低的計算復雜度，證明可作為一般醫(yī)學圖像分割的一種替代網(wǎng)絡。

由于使用的預訓練權(quán)重并不是從端到端的訓練生成的，一定程度上影響了網(wǎng)絡的性能。下一步工作是解決預訓練模型的問題，獲得更好的預訓練權(quán)重。其次就是繼續(xù)優(yōu)化模型結(jié)構(gòu)，將CoT-TransUNet推廣到3D醫(yī)學圖像分割中。