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