各位做醫(yī)學(xué)圖像分割的朋友今天聊一個繞不開的模型——TransUNet。如果你手頭正好有CT、MRI或者病理切片的分割需求又覺得純CNN的U-Net雖然好但總是差點全局語義信息那這篇文章應(yīng)該能給你一些啟發(fā)。TransUNet是2021年發(fā)表在NeurIPS上的工作核心就一句話把Transformer當(dāng)成編碼器的增強組件嵌入到U-Net框架里用CNN提取底層空間特征用Transformer建模長距離依賴最后通過解碼器逐步恢復(fù)分辨率。它在多器官分割、心臟分割等任務(wù)上的表現(xiàn)確實證明了這套混合架構(gòu)的價值。這篇博文不打算只復(fù)述論文我會結(jié)合自己復(fù)現(xiàn)和改造TransUNet的實際經(jīng)驗把它的設(shè)計邏輯、實現(xiàn)細(xì)節(jié)、訓(xùn)練技巧和容易踩的坑都攤開來講。不管你是剛?cè)腴T的學(xué)生還是已經(jīng)在做落地項目的工程師只要你用得上醫(yī)學(xué)圖像分割這篇文章應(yīng)該都能給你一些實實在在的參考。1. 為什么偏偏是TransUNetU-Net的邊界與Transformer的機會1.1 純U-Net在醫(yī)學(xué)分割里的瓶頸在深入TransUNet之前得先搞清楚它要解決的問題是什么。U-Net從2015年提出至今一直是醫(yī)學(xué)圖像分割的主力軍原因很簡單它在標(biāo)注數(shù)據(jù)極其有限的情況下依然能靠數(shù)據(jù)增強和編碼-解碼結(jié)構(gòu)學(xué)出不錯的結(jié)果。但用過U-Net的人應(yīng)該都有體會隨著層數(shù)加深特征圖在不斷下采樣的過程中會丟失不少空間細(xì)節(jié)。雖然跳躍連接能帶來一些補償?shù)捎诿總€卷積核的感受野始終有限U-Net很難對相距很遠(yuǎn)的像素之間建立關(guān)系。舉個例子在肝臟分割任務(wù)里肝臟區(qū)域的邊緣通常和周圍組織灰度對比并不強烈機器很容易把相鄰的胃壁或肌肉組織誤認(rèn)成肝臟。如果網(wǎng)絡(luò)只盯著局部紋理就非常容易產(chǎn)生這種邊界混淆。想讓模型真正理解“這是一個完整的肝臟”就必須讓它對整張影像有一個全局級別的感知。傳統(tǒng)的做法是把網(wǎng)絡(luò)加深、加寬或者引入空洞卷積這雖然能擴(kuò)大感受野但計算開銷上去了很多小顯存的卡根本跑不動。1.2 Transformer帶來的可能性與它的尷尬Transformer在自然語言處理領(lǐng)域的成功讓很多人看到了它在視覺任務(wù)上的潛力。它的自注意力機制能在一層之內(nèi)讓任意兩個位置直接交互也就是說不管兩個像素在空間上隔得多遠(yuǎn)模型都有能力學(xué)習(xí)它們之間的關(guān)聯(lián)。這一點恰恰彌補了CNN的短板。但這里也有一個尷尬的問題Transformer在處理高分辨率圖像時計算量是隨序列長度的平方增長的。醫(yī)學(xué)圖像通常都是512x512、甚至1024x1024級別的單通道或三通道圖如果你直接把每個像素當(dāng)成一個token送進(jìn)Transformer那計算量是任何一張顯卡都承受不起的。所以ViT的做法是把圖像切分成16x16的小塊每一塊當(dāng)成一個token。但這樣一來想要做到像素級精細(xì)分割Transformer對空間細(xì)節(jié)的恢復(fù)能力就不夠了。1.3 TransUNet的設(shè)計思路不是二選一而是混合TransUNet的核心聰明之處在于它沒有在CNN和Transformer之間做非此即彼的選擇而是把兩者組合在一起。具體來說先用一個CNN作為骨干網(wǎng)絡(luò)提取特征得到一層分辨率適中、語義豐富的特征圖然后把這張?zhí)卣鲌D切成一個個patch送入Transformer編碼器去建模全局關(guān)系最后把經(jīng)過Transformer處理的特征重新拼成二維特征圖通過U-Net風(fēng)格的解碼器逐級恢復(fù)分辨率。這個設(shè)計的巧妙之處在于CNN部分負(fù)責(zé)搞定細(xì)粒度的空間結(jié)構(gòu)Transformer部分負(fù)責(zé)搞定全局語義而解碼器部分利用跳躍連接把底層細(xì)節(jié)和高層語義融合起來。我可以直接說這個組合在多個醫(yī)學(xué)分割基準(zhǔn)上的確比單純的U-Net、單純的TransUNet變體以及很多CNN與Transformer簡單拼接的模型效果都要好。它并不是一個理論完美、但實操很難受的模型相反它對顯存的要求和訓(xùn)練時間都處在可以接受的范圍這也是它能被廣泛復(fù)現(xiàn)和使用的原因。2. 深入網(wǎng)絡(luò)骨架編碼器、解碼器與跳躍連接的設(shè)計考量2.1 CNN編碼器用什么骨干網(wǎng)絡(luò)決定了特征質(zhì)量TransUNet整篇論文最容易被忽視、但實操中影響最大的部分其實是Encoder的CNN骨干網(wǎng)絡(luò)選擇。論文里實驗了兩種主流選擇一種是經(jīng)典的ResNet50另一種是輕量級的MobileNetV2。作者在不同實驗里用了不同的backbone來驗證他們的方法。我自己跑下來的感受是選擇ResNet50作為骨干時模型整體的分割精度會比輕量級網(wǎng)絡(luò)高一截但訓(xùn)練時間和顯存占用都會明顯增加。如果你在資源有限的場景下工作MobileNetV2版本的TransUNet仍然有不錯的可用性尤其適合快速驗證和部署。MobileNetV2靠深度可分離卷積大幅減少了參數(shù)量在差不多精度的情況下模型大小可能只有ResNet50版本的一半都不止。我的建議是如果你面向的是落地推理、邊緣設(shè)備部署那優(yōu)先選MobileNetV2版本如果你做的是學(xué)術(shù)實驗、目標(biāo)是在排行榜上刷精度那ResNet50是更穩(wěn)妥的選擇。2.2 Transformer編碼器Patch嵌入和自注意力的細(xì)節(jié)Transformer編碼器部分的輸入來自CNN骨干網(wǎng)絡(luò)輸出的特征圖。比如輸入圖像是224x224經(jīng)過ResNet50多次下采樣后我們拿到的是14x14的空間尺寸、通道數(shù)為512或768的特征圖。接著要將這個特征圖切分成一個個patch并用一個線性層做patch嵌入。14x14的特征圖可以被切成1x1的patch也可以切成2x2的patch。如果patch大小越大得到的token數(shù)量就越少計算量越小但局部信息的粒度就越粗。論文里大量實驗使用1x1的patch因為此時14x14等于196個token計算量相當(dāng)小而且能在保持較高分辨率的同時讓Transformer捕捉全局關(guān)系。你可以在自己的實驗中按需調(diào)節(jié)這個參數(shù)但我個人的經(jīng)驗是醫(yī)學(xué)圖像的病灶結(jié)構(gòu)通常比較緊湊1x1的patch輸出效果往往最穩(wěn)。Transformer層里還有兩個關(guān)鍵的細(xì)節(jié)位置編碼和類別token。位置編碼用來告訴模型每個token在空間上的相對位置如果缺了這個信息Transformer就很難區(qū)分兩張內(nèi)容相同但位置不同的特征圖。論文使用的是標(biāo)準(zhǔn)可學(xué)習(xí)位置編碼訓(xùn)練時直接作為參數(shù)優(yōu)化。而類別token是ViT里用來做圖像分類的在TransUNet里因為任務(wù)不是分類這個token最終會被丟棄或掩碼掉重點放在特征重組部分。2.3 解碼器和跳躍連接空間分辨率的恢復(fù)路徑經(jīng)過Transformer編碼器輸出的特征序列需要重新reshape回二維特征圖比如從196個token恢復(fù)成14x14的空間形狀。然后進(jìn)入解碼器。TransUNet的解碼器設(shè)計大體沿用U-Net的結(jié)構(gòu)每一級先通過上采樣操作放大分辨率再將來自編碼器同層級的特征拼接進(jìn)來融合后經(jīng)過卷積層精細(xì)調(diào)整。如此反復(fù)直到恢復(fù)成和原始輸入相同的分辨率。之所以要拼接編碼器特征是因為高層特征雖然有豐富語義但空間位置是很粗糙的底層特征雖然語義弱但邊緣和紋理信息保留得很好。兩者結(jié)合網(wǎng)絡(luò)既能知道“這塊區(qū)域是個器官”又能定位“這個器官的邊界在哪個像素”。這個設(shè)計的另一個好處是緩解了Transformer帶來的局部細(xì)節(jié)損失。Transformer在建模全局關(guān)系時并不會天然保留像素級的邊界信息跳躍連接等于在恢復(fù)階段把這一層信息從底層直接引回來避免解碼器只靠高層特征而丟失邊緣。3. 從公式到代碼損失函數(shù)、訓(xùn)練配置與復(fù)現(xiàn)實驗3.1 輸入尺寸與預(yù)處理并非越大越好先把最實際的訓(xùn)練配置講清楚。醫(yī)學(xué)圖像分割的輸入尺寸需要根據(jù)GPU顯存情況去權(quán)衡。我最初嘗試的時候直接用512x512的分辨率輸入期望保留盡可能多的細(xì)節(jié)結(jié)果一個batch只能放兩張圖訓(xùn)練速度慢得讓人崩潰而且小目標(biāo)組織分割的精度并沒有明顯比224x224高很多。后來我做了實驗對比在Synapse多器官分割數(shù)據(jù)集上224x224輸入配合resize和隨機裁剪分割效果和384x384輸入非常接近但訓(xùn)練時間是后者的三分之一不到。原因在于多器官分割任務(wù)中各個器官的空間分布和邊界在較低分辨率下其實已經(jīng)能夠被較好估計更關(guān)鍵的是語義信息而不是微觀紋理。如果你的顯存只有12GB左右建議優(yōu)先使用224x224并使用一些在線數(shù)據(jù)增強來彌補信息損失如果你有24GB以上的顯存可以直接上384x384甚至更高。預(yù)處理方面醫(yī)學(xué)圖像的窗寬窗位調(diào)整至關(guān)重要。以CT影像為例不同組織的CT值范圍差異很大。一般建議將圖像裁剪到感興趣的窗寬范圍內(nèi)比如肝臟分割常在[-150, 250]之間再把像素歸一化到[0, 1]或均值為0、方差為1。沒有這一步模型會花費很多不必要的參數(shù)去區(qū)分不同患者之間的掃描差異。3.2 損失函數(shù)選擇交叉熵、Dice Loss還是混合損失分類任務(wù)里大家習(xí)慣用交叉熵但醫(yī)學(xué)分割任務(wù)有一個非常典型的問題前景和背景的像素數(shù)量極度不平衡。比如一個512x512的CT切片里肝臟可能只占幾百個人像素背景卻有幾十萬個。這時如果只用交叉熵模型很容易把所有像素都預(yù)測為背景來換取很低的loss。常見的解決方案是Dice Loss它直接優(yōu)化分割結(jié)果和標(biāo)注之間的重合度。它的公式很簡單等于2乘以預(yù)測和標(biāo)注的交集除以上兩者的并集。Dice Loss對類別不平衡問題有很強的容忍性因為不管前景多小只要沒預(yù)測準(zhǔn)loss就會很高。但我實際用下來純Dice Loss在訓(xùn)練初期非常不穩(wěn)定因為梯度變化太劇烈導(dǎo)致網(wǎng)絡(luò)收斂困難。更穩(wěn)的做法是把Dice Loss和交叉熵組合起來。我常用的是Dice Loss 加權(quán)交叉熵權(quán)重分別設(shè)為0.5和0.5。這樣既保持了交叉熵在像素級別上的平滑梯度又能利用Dice Loss加強對小目標(biāo)的約束。在Synapse數(shù)據(jù)集的實驗中這種組合比單獨使用Dice Loss提高了約1.5%的Dice系數(shù)。3.3 訓(xùn)練策略從預(yù)訓(xùn)練權(quán)重開始的收益TransUNet的Transformer部分建議從預(yù)訓(xùn)練權(quán)重初始化。這個預(yù)訓(xùn)練通常是在ImageNet上用圖像分類任務(wù)訓(xùn)練的ViT或DeiT權(quán)重。有些人可能會想醫(yī)學(xué)圖像和自然圖像差異這么大直接用預(yù)訓(xùn)練權(quán)重會不會反而不好我的實測結(jié)果是即使考慮了這個差異從預(yù)訓(xùn)練權(quán)重初始化的模型在分割精度上仍然要明顯優(yōu)于隨機初始化并且收斂速度快得多。原因也好理解Transformer的底層特征學(xué)習(xí)的是通用的視覺結(jié)構(gòu)比如邊緣、顏色、紋理組合這些特征在自然圖像和醫(yī)學(xué)圖像之間是相通的。用預(yù)訓(xùn)練權(quán)重相當(dāng)于讓模型站在了巨人的肩膀上只花少量數(shù)據(jù)去適應(yīng)醫(yī)學(xué)圖像的分布即可。建議在訓(xùn)練前先凍結(jié)骨干網(wǎng)絡(luò)和Transformer編碼器的前幾層只改動解碼器部分觀察loss變化待整個模型有了一個基礎(chǔ)擬合能力后再全部解凍這樣能有效防止訓(xùn)練初期的大幅震蕩。3.4 實驗復(fù)現(xiàn)一個可以用作baseline的訓(xùn)練流程示例我自己復(fù)現(xiàn)時使用的訓(xùn)練框架是PyTorch配合兩個GPUbatch size設(shè)為16初始學(xué)習(xí)率設(shè)成了1e-4采用AdamW優(yōu)化器設(shè)置weight decay為1e-4。整個流程可以參考下面的偽代碼結(jié)構(gòu)model TransUNet( img_size224, in_channels1, num_classesnum_classes, embed_dim512, depth12, num_heads8, patch_size1, backboneresnet50 ) criterion CombinedLoss(ce_weight0.5, dice_weight0.5) optimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max200) for epoch in range(200): for images, labels in train_loader: images, labels images.cuda(), labels.cuda() outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() scheduler.step() # 關(guān)于驗證集上的Dice和IoU指標(biāo)統(tǒng)計 eval_model(model, val_loader)訓(xùn)練過程中我建議每5個epoch保存一次checkpoint同時保留一個專門的best模型。很多情況下模型在后期會出現(xiàn)過擬合表現(xiàn)在驗證集Dice不再上升但訓(xùn)練集loss還在下降。如果遇到這種情況果斷回退到驗證指標(biāo)最好的那個權(quán)重或者提前終止訓(xùn)練。Synapse數(shù)據(jù)集上我這個配置大約在120個epoch左右就能達(dá)到一個穩(wěn)定的分割效果。4. 從零到一實現(xiàn)時的硬核細(xì)節(jié)與性能優(yōu)化4.1 訓(xùn)練內(nèi)存優(yōu)化混合精度與梯度累積醫(yī)學(xué)圖像分割模型訓(xùn)練內(nèi)存占用高這是不少人勸退的原因。以TransUNet為例輸入224x224、batch size為16的時候一張24GB的顯卡可能會被擠到只剩幾個G。如果batch size只能設(shè)為2或4那么模型訓(xùn)練的效果會明顯變差因為Batch Normalization層的統(tǒng)計信息在很小的batch里會非常不穩(wěn)定。一個直接的解決辦法是使用梯度累積。簡單說就是把若干個batch的梯度累加起來達(dá)到目標(biāo)batch size對應(yīng)的梯度后再執(zhí)行一次參數(shù)更新。這樣我們雖然每次只喂少量圖像進(jìn)模型但梯度的累積效果跟完整batch一樣。另一個手段是開啟PyTorch自帶的混合精度訓(xùn)練也就是AMP。開啟AMP后模型在前向和反向傳播時自動將部分計算切換到FP16半精度顯存占用能降一半左右訓(xùn)練速度還能提升不少。我在啟用AMP和梯度累積后原來需要24GB顯存才能跑起來的配置在12GB顯卡上也順利完成了訓(xùn)練。4.2 數(shù)據(jù)加載效率別讓IO成為訓(xùn)練瓶頸很多人在模型上花了很多心思卻忽略了數(shù)據(jù)加載這一環(huán)。尤其是在醫(yī)學(xué)圖像場景每次讀取的都是NIfTI格式的三維數(shù)據(jù)或者整張全切片IO開銷相當(dāng)大。我最初用普通DataLoader時數(shù)據(jù)加載時間居然占了總訓(xùn)練時間的40%。需要針對性地做一些優(yōu)化。首先是實現(xiàn)一個高效的數(shù)據(jù)加載類在__getitem__函數(shù)里做切片和預(yù)處理時盡量用numpy的切片操作而不是循環(huán)逐像素處理。其次是開啟DataLoader的num_workers參數(shù)通常設(shè)置為CPU核心數(shù)的一半數(shù)據(jù)加載就能顯著提速。另外如果內(nèi)存足夠可以一次把整個訓(xùn)練集讀入內(nèi)存在內(nèi)存里完成切片和數(shù)據(jù)增強這是最省時的方案。對于超大切片建議預(yù)先做patches的切割把每張訓(xùn)練圖像切成若干個固定大小的patch保存下來訓(xùn)練時直接讀取patch文件。4.3 推理階段加速減少冗余計算的思路訓(xùn)練結(jié)束后推理階段的效率也很關(guān)鍵。TransUNet中Transformer編碼器部分是計算量最大的模塊之一。如果你要在多張圖像上做推理可以嘗試把Transformer編碼器在訓(xùn)練時學(xué)習(xí)的權(quán)重固定不動只對圖像做一次特征提取和一次全局建模。對于同一組參數(shù)不需要在每張測試圖上重新計算前向過程中的某些臨時變量PyTorch的torch.no_grad()模式可以節(jié)省大量不必要的顯存和計算時間。另外一個更實際的加速策略是模型剪枝和蒸餾。可以把TransUNet當(dāng)成教師模型訓(xùn)練一個結(jié)構(gòu)更簡單的學(xué)生模型來模擬教師模型的輸出。這種方法在醫(yī)學(xué)圖像場景已經(jīng)有了一些成功案例學(xué)生模型的速度能提升數(shù)倍同時精度損失控制在可接受范圍內(nèi)。不過這塊內(nèi)容比較深需要分?jǐn)?shù)據(jù)去實驗適合已經(jīng)對TransUNet本身有充分理解的讀者。5. 常見問題與排查技巧實錄5.1 訓(xùn)練不收斂先檢查這五件事如果你用TransUNet訓(xùn)練時發(fā)現(xiàn)loss幾乎不下降或者Dice系數(shù)一直在零點幾徘徊先別急著調(diào)模型結(jié)構(gòu)按順序檢查下面五件事第一檢查數(shù)據(jù)歸一化是否合理。醫(yī)學(xué)圖像如果直接用原始像素值送入模型數(shù)值范圍可能非常大導(dǎo)致梯度爆炸。CT圖像尤其要注意窗寬窗位處理。第二檢查標(biāo)簽編碼。有些醫(yī)學(xué)數(shù)據(jù)集的標(biāo)注文件用的是255表示目標(biāo)區(qū)域、0表示背景如果直接用交叉熵計算模型很難收斂正確的做法是把標(biāo)簽改成類別索引比如0、1、2。第三檢查損失函數(shù)是否寫對。Dice Loss的公式里分子分母如果計算維度不對很容易出現(xiàn)范圍內(nèi)浮點誤差導(dǎo)致的NaN。第四檢查學(xué)習(xí)率是否合適。Transformer對學(xué)習(xí)率比較敏感1e-4是個相對穩(wěn)妥的起點如果你用了更大的學(xué)習(xí)率建議先降到3e-5再試試。最后檢查類別權(quán)重是否設(shè)置正確。如果你對多器官分割任務(wù)按器官類別設(shè)置權(quán)重權(quán)重差距過大也會導(dǎo)致訓(xùn)練不穩(wěn)定。5.2 分割結(jié)果出現(xiàn)空洞或邊緣毛刺這類問題相信很多做分割的人都遇到過。模型輸出的結(jié)果有時候會呈現(xiàn)不規(guī)則的雜點或者在器官內(nèi)部出現(xiàn)被誤分割成背景的小洞。這通常是因為模型在解碼器的最后一層上特征圖分辨率不夠細(xì)膩加上損失函數(shù)對邊界像素的約束不夠嚴(yán)格。我的處理經(jīng)驗是在訓(xùn)練階段增加邊界感知的損失項。具體來說可以對標(biāo)注圖像做一次邊緣提取把邊緣像素的權(quán)重調(diào)得更高讓模型更關(guān)注這些困難位置。還可以在最后輸出層后增加一個條件隨機場后處理能夠有效整理分割邊緣。不過CRF的推理速度偏慢如果是在線服務(wù)場景需要權(quán)衡。最直接的辦法是調(diào)整patch大小和輸入分辨率往往能把邊緣質(zhì)量提升一個檔次。5.3 多個器官互相粘連邊界區(qū)分困難Synapse這類多器官數(shù)據(jù)集上肝臟和脾臟、胃和胰腺之間的距離非常近灰度分布甚至很接近模型經(jīng)常把它們連在一起。前期我跑的時候器官粘連問題特別突出測評的Dice指標(biāo)倒還可以但看了一眼可視化結(jié)果發(fā)現(xiàn)兩個器官完全糊成一片。解決辦法之一是在標(biāo)簽處理上利用器官之間的相對位置信息。Transformer的一個優(yōu)勢就是能捕獲這種遠(yuǎn)距離關(guān)系所以如果你發(fā)現(xiàn)模型還是粘連很有可能是Transformer層數(shù)不夠深或者patch尺寸太大導(dǎo)致分割很粗糙?？梢試L試加深深度比如從12層編碼器變成16層并把patch從2x2改成1x1。我在實際實驗中發(fā)現(xiàn)將patch設(shè)為1x1同時加深兩層Synapse數(shù)據(jù)集上的臟器分離效果就有明顯改善。6. 適用場景與擴(kuò)展方向哪些任務(wù)適合用它哪些未必6.1 適合TransUNet的任務(wù)特點我總結(jié)了一下TransUNet特別適合那些既需要精細(xì)邊界、又需要全局語義的任務(wù)。多器官分割、心臟結(jié)構(gòu)分割、腫瘤區(qū)域分割、視網(wǎng)膜血管分割這些任務(wù)要么關(guān)注的目標(biāo)尺度變化很大要么目標(biāo)分布不太規(guī)則需要全局上下文輔助判斷。TransUNet在這些場景下通常能比U-Net高出幾個百分點的Dice得分。具體來說在心臟MRI分割任務(wù)中心肌的形狀相對穩(wěn)定但心腔的邊緣和周圍組織對比度低這時Transformer的全局感知能幫助減少誤判。在肝臟和肝腫瘤分割任務(wù)中腫瘤的尺寸、形狀變化極大純CNN容易把零散小結(jié)節(jié)漏掉TransUNet的全局建模則能捕捉到腫瘤和非腫瘤組織的微妙差異。6.2 不一定適合的任務(wù)你該怎么判斷TransUNet也不是萬能的對于那種圖像尺寸巨大、目標(biāo)非常小、且目標(biāo)之間沒有明顯空間關(guān)聯(lián)的任務(wù)它的優(yōu)勢不一定明顯但計算開銷卻很大。比如血管分割血管形態(tài)細(xì)長且在整個圖像中分布稀疏Transformer雖然能覆蓋全局但特征過粗時反而容易忽略細(xì)小分支。另一個不太適合的是實時推理場景。因為Transformer編碼器的計算量畢竟比純卷積大不少如果對單張圖像的推理時間要求在毫秒級TransUNet帶來的提升可能不值得犧牲的延遲。對這些任務(wù)可以考慮輕量化的變體比如把TransUNet編碼器替換成更小的Tiny版本或者采用知識蒸餾的方式只學(xué)習(xí)主流器官的信息、降低解碼器的復(fù)雜度。6.3 衍生方向TransUNet之后它還留下了哪些思路TransUNet本身只是一個開始它啟發(fā)了大量后續(xù)工作這些工作基本都是在它的框架上做了局部創(chuàng)新。有的把CNN骨干替換成了更強大的ConvNeXt或Swin Transformer有的把Transformer編碼器換成了Swin Transformer因為Swin使用窗口注意力機制計算效率和局部建模能力更強也有的引入了多尺度的思想通過不同層級的特征圖分別計算Transformer attention再融合結(jié)果。如果你想在這個方向繼續(xù)深入建議先吃透TransUNet的代碼實現(xiàn)再去理解它的衍生模型會節(jié)省很多周折。我個人在實際操作中的體會是做醫(yī)學(xué)圖像分割拼模型結(jié)構(gòu)是最簡單的一步真正花時間的是數(shù)據(jù)處理、損失函數(shù)設(shè)計和調(diào)參。TransUNet提供了很好的骨架但最終效果的差異往往取決于你怎么處理數(shù)據(jù)、怎么設(shè)計訓(xùn)練流程、怎么解決落地場景里的各種臟問題。如果準(zhǔn)備跑這個模型可以先拿Synapse或ACDC數(shù)據(jù)集跑通整個流程再遷移到自己的數(shù)據(jù)上這個路徑相對平緩能讓人少走很多彎路。