CoT-TransUNet:輕量化的上下文Transformer醫(yī)學(xué)圖像分割網(wǎng)絡(luò)

楊 鶴，柏正堯

云南大學(xué) 信息學(xué)院，昆明 650500

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

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

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

1 相關(guān)工作

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

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

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

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

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

1.3 CARAFE上采樣算子

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

2 方法

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

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

2.1 混合編碼器

2.1.1 特征提取

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

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

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

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

2.1.2 圖像序列化

Transformer塊：對(duì)于一張圖像X∈ΦH×W×C，其空間分辨率為H×W，通道數(shù)為C。將輸入X重構(gòu)為一些2D的，其中每個(gè)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為輸入生成特征映射。再對(duì)從CNN特征圖中提取的1×1個(gè)patch進(jìn)行嵌入，而不是對(duì)原始圖像進(jìn)行嵌入[8]。

2.2 解碼器

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

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

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

2.3 評(píng)價(jià)指標(biāo)

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

2.3.1 Dice系數(shù)

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

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

2.3.2 Hausdorff距離

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

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

3 實(shí)驗(yàn)

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

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

3.2 實(shí)施細(xì)節(jié)

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

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

3.3 實(shí)驗(yàn)結(jié)果

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

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

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

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

3.4 分析研究

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

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

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

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

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

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

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

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

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

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

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

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

3.5 討論

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

4 結(jié)論

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

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