TransBTS_3D多模態(tài)腦腫瘤分割Transformer閱讀筆記

文章目錄

• 1 前言
• 2 TransBTS 整體結(jié)構(gòu)概述
• 3 Network Encoder
• 4 Network Decoder
• 5 實驗部分
• 6 總結(jié)
• 7 參考鏈接

1 前言

這是醫(yī)學(xué)圖像處理系列的第三篇文章，arXiv 鏈接在文末，為什么選擇寫下 TransBTS 這篇論文的閱讀筆記？因為才有時間整理了。此外，與之前的兩篇相比，它們都有各自的代表性。第一篇（https://zhuanlan.zhihu.com/p/505483978）是 2D Transformer 醫(yī)學(xué)圖像分割，第二篇（https://zhuanlan.zhihu.com/p/506716802）是輕量級網(wǎng)絡(luò) 2.5D 的，在這一篇閱讀筆記中，講解的則是 3D 多模態(tài)腦腫瘤分割 Transformer，也接受在 MICCAI 上。后面文章過多的話，我會整理成目錄的形式。

Transformer 由于自注意力機(jī)制具有捕獲全局依賴特征的能力，且允許網(wǎng)絡(luò)根據(jù)輸入內(nèi)容動態(tài)收集相關(guān)特征，在自然語言處理和 2D 圖像分類方面取得了成功。

然而，局部特征和全局特征對于 dense 的預(yù)測任務(wù)同樣重要，尤其是對于 3D 醫(yī)學(xué)圖像分割。那么在 3D CNN 中利用 Transformer 是值得關(guān)注的，這篇文章中提出了基于編碼器-解碼器結(jié)構(gòu)的新型網(wǎng)絡(luò) TransBTS，BTS 指的是 MRI 的腦腫瘤分割。

為了獲取局部的三維上下文信息，編碼器首先利用 3D CNN 提取三維空間特征圖。同時，對特征映射進(jìn)行了細(xì)致的改造，將處理后的向量（token）輸入到 Transformer 中進(jìn)行全局特征建模。解碼器利用 Transformer 嵌入的特征，進(jìn)行漸進(jìn)式上采樣來預(yù)測詳細(xì)的分割圖。

本篇閱讀筆記首先對論文進(jìn)行概述；再詳細(xì)介紹論文提出的 TransBTS 結(jié)構(gòu)，設(shè)計一個神經(jīng)網(wǎng)絡(luò)，通過利用 Transformer 來有效地對 3D 醫(yī)學(xué)圖像的空間和深度維度上的局部和全局特征進(jìn)行建模；最后分析了實驗和分割可視化部分，并簡單的總結(jié)。

2 TransBTS 整體結(jié)構(gòu)概述

這部分我們直接看 TransBTS 的整體結(jié)構(gòu)，然后再分析具體細(xì)節(jié)（分別看 Encoder 和Decoder）。關(guān)于 Transformer 在醫(yī)學(xué)圖像上應(yīng)用的細(xì)節(jié)，可以在第一篇閱讀筆記中發(fā)現(xiàn)。

對于給定輸入圖像 X，尺寸為 C×H×W×D，空間分辨率為H×W，深度維度為 D（切片數(shù)）和 C（模態(tài)數(shù)），首先利用 3D CNN 生成了捕獲空間和深度信息的緊湊特征圖（下采樣），然后利用 Transformer layer 在全局空間中對長距離相關(guān)特征進(jìn)行建模。

之后，重復(fù)疊加上采樣和卷積層，以逐漸產(chǎn)生高分辨率的分割結(jié)果，因為，高分辨率是醫(yī)學(xué)圖像必須的特性。如下圖，這里我們先不關(guān)注具體細(xì)節(jié)，網(wǎng)絡(luò)結(jié)構(gòu)還是很清晰的。

3 Network Encoder

這部分我們先看一下網(wǎng)絡(luò)的編碼器部分，有一些細(xì)節(jié)需要介紹。由于 Transformer 的計算復(fù)雜度是 token 數(shù)量（即序列長度）的二次方，因此直接將輸入圖像展平為作為 Transformer 輸入的序列是不切實際的。

因此，ViT 將圖像拆分為固定大小 16 × 16 的 patch，然后將每個 patch 作為一個 token，將序列長度減少到 16 的平方。同理，對于 3D 圖像數(shù)據(jù)，則轉(zhuǎn)換為 3D patch 就好（為了方便下面理解，也就是說輸入的特征圖多了深度）。

然而，上面的策略使得 Transformer 無法跨空間和深度維度對圖像局部上下文信息進(jìn)行建模以進(jìn)行 3D 分割（局部 3D 上下文特征獲取難）。

為了解決這個關(guān)鍵問題，這篇論文中的解決方案是使用下采樣（stride=2 的卷積）堆疊 3 × 3 × 3 卷積塊，以逐漸將輸入圖像編碼為低分辨率的高級特征表示 F，尺寸為 K × H/8 × W/8 ×D/8 (K = 128)，既 H、W 和 D 是輸入維度的 1/8（總步長 = 8）。

通過這種方式，豐富的局部 3D 上下文特征有效地嵌入到 F 中。然后，將 F 輸入到 Transformer 編碼器中，以進(jìn)一步學(xué)習(xí)全局感受野的相關(guān)特征。

關(guān)于論文中的 Feature Embedding of Transformer Encoder。給定上面提到的特征圖 F，為了確保每個 volume 的全面表示，使用線性投影（一個 3×3×3 卷積層）將通道維度從 K = 128 增加到 d = 512。Transformer 層期望一個序列作為輸入。

因此，需要將空間和深度維度折疊為一維，得到一個 d×N（N = H/8 × W/8 × D/8）特征圖 f ，也可以看作是 N 個 size 為 d 的 token。

此外，為了對分割任務(wù)中重要的位置信息進(jìn)行編碼，使用了可學(xué)習(xí)的 position embedding，并直接將它們與特征圖 f 相加，這個操作在上面的整體網(wǎng)絡(luò)結(jié)構(gòu)圖中可以看到哈，我就不贅述了。

然后就輸入到 Transformer Layers 了，用 Transformer 的一些規(guī)則計算（同樣可以在第一篇閱讀筆記中發(fā)現(xiàn)），在這篇閱讀筆記中，我也不詳細(xì)介紹了。至此，就是 TransBTS 的編碼器部分。

4 Network Decoder

下面我們來看下 TransBTS 的解碼器部分。MRI 腦腫瘤分割需要在原始 3D 圖像空間（H×W×D）中生成分割結(jié)果，所以引入了 3D CNN 解碼器來執(zhí)行特征上采樣和像素級分割，也就是第二部分網(wǎng)絡(luò)結(jié)構(gòu)中的右半部分。

我們注意到，在右半部分大概有兩個值得關(guān)注的內(nèi)容，分別為 Feature Mapping 和 Progressive Feature Upsampling。

關(guān)于 Feature Mapping，它在解碼器的首部，為了適應(yīng) 3D CNN 解碼器的輸入維度，特征映射模塊將序列數(shù)據(jù)投影回標(biāo)準(zhǔn)的 4D 特征映射。此外，為了降低解碼器的計算復(fù)雜度，又采用卷積塊將通道維度從 d 減小到 K。那么現(xiàn)在特征圖的恢復(fù)到 K × H/8 × W/8 × D/8。

繼續(xù)往下看，還有 Progressive Feature Upsampling 操作。也就是通過一系列的上采樣操作和卷積塊，得到 H × W × D 的分割結(jié)果。此外，還有跳過連接，以獲得更精細(xì)的分割結(jié)果和更豐富的空間細(xì)節(jié)。解碼器相對簡單，到這就結(jié)束了。

5 實驗部分

在這一部分，首先來看下這個 3D 多模態(tài)數(shù)據(jù)集。數(shù)據(jù)集中每個樣本由四種腦磁共振成像掃描模式組成，即 T1、T1ce、T2 和 FLAIR。每個圖片都是 240 × 240×155，并排列在同一個空間中。

標(biāo)簽包含 4 類：背景（標(biāo)記 0）、壞死和非增強(qiáng)腫瘤（標(biāo)記 1）、水腫（標(biāo)記 2）和增強(qiáng)腫瘤（標(biāo)記 4）。分割精度通過用于增強(qiáng)腫瘤區(qū)域（ET，標(biāo)記 1）、腫瘤核心區(qū)域（TC，標(biāo)記 1 和 4）和整個腫瘤區(qū)域（WT，標(biāo)記 1、2 和 4）的 Dice 評分和 Hausdorff 距離（95%）度量來測量。

實驗表明，TransBTS 在 3D MRI 掃描分割上實現(xiàn)了優(yōu)于最先進(jìn)的腦腫瘤分割方法的性能，對比實驗結(jié)果如下。

相應(yīng)的，對結(jié)果進(jìn)行可視化，如下圖。可以明顯看出，TransBTS 可以更準(zhǔn)確地分割出腦腫瘤，通過對每個體積之間的相關(guān)性進(jìn)行建模來生成更好的分割 mask。

最后，我們來看下消融實驗的部分。分別為序列長度消融實驗、Transformer 位置的消融實驗（下采樣四次效果最好）和跳躍連接位置的消融研究，如下所示。

6 總結(jié)

這篇閱讀筆記介紹了一種新的醫(yī)學(xué)圖像分割框架，該框架有效地結(jié)合了 3D CNN 和Transformer，用于磁共振成像中的多模態(tài)腦腫瘤分割。TransBTS 不僅繼承了 3D CNN 建模局部上下文信息的優(yōu)勢，而且利用了 Transformer 學(xué)習(xí)全局語義相關(guān)性。最后，在 BraTS 2019 數(shù)據(jù)集上的實驗結(jié)果驗證了所提出的 TransBTS 的有效性。

這是醫(yī)學(xué)圖像處理系列的第三篇閱讀筆記，歡迎交流和糾正！

7 參考鏈接

• https://arxiv.org/pdf/2103.04430
• https://github.com/Wenxuan-1119/TransBTS