1. 從一張病理切片說起為什么傳統(tǒng)卷積網(wǎng)絡(luò)在醫(yī)學(xué)圖像上總差那么一口氣第一次接觸醫(yī)學(xué)圖像分割是幫一個(gè)做病理分析的朋友處理一批肝臟切片。那批數(shù)據(jù)里腫瘤區(qū)域和正常組織的邊界模糊得像水墨畫暈開的邊緣用當(dāng)時(shí)主流的卷積網(wǎng)絡(luò)跑出來的掩膜邊緣總是毛毛躁躁要么把正常組織誤吞進(jìn)去要么在腫瘤內(nèi)部留下空洞。調(diào)了各種損失函數(shù)、加了后處理效果提升始終有限。后來才意識到問題不在訓(xùn)練技巧而在網(wǎng)絡(luò)本身的視野——卷積核的感受野是局部的、固定的它天生不擅長處理那種整體形狀依賴局部細(xì)節(jié)、局部細(xì)節(jié)又反過來定義整體的醫(yī)學(xué)圖像。醫(yī)學(xué)圖像有個(gè)很討厭的特點(diǎn)目標(biāo)邊界往往不是靠單一尺度的紋理決定的。比如肺結(jié)節(jié)它的良惡性判斷既依賴結(jié)節(jié)整體的圓形度、邊緣光滑度也依賴內(nèi)部鈣化點(diǎn)的細(xì)微紋理再比如視網(wǎng)膜血管分割粗血管的走向和毛細(xì)血管的末梢是兩種完全不同的尺度信息但它們在解剖上是連續(xù)的。卷積網(wǎng)絡(luò)堆疊再深本質(zhì)上還是在做局部特征的逐層抽象跨尺度的全局關(guān)聯(lián)需要靠很深的網(wǎng)絡(luò)去間接建立效率低且容易丟失信息。這就是TransNeXt這類架構(gòu)切入的痛點(diǎn)。它把Transformer的全局建模能力和卷積的局部歸納偏置揉在一起核心武器就是標(biāo)題里提到的聚合注意力Aggregated Attention。這個(gè)機(jī)制的設(shè)計(jì)靈感說白了就是模仿人眼看東西的方式——你不會(huì)盯著一個(gè)像素一個(gè)像素地看而是先用余光掃全局輪廓再把注意力聚焦到關(guān)鍵區(qū)域同時(shí)大腦還在不斷把這兩種信息聚合起來。TransNeXt想做的就是讓網(wǎng)絡(luò)也具備這種全局掃視局部聚焦動(dòng)態(tài)聚合的能力。這篇文章我會(huì)從實(shí)際落地的角度把TransNeXt在醫(yī)學(xué)圖像分割任務(wù)里的設(shè)計(jì)思路、聚合注意力的原理、完整的實(shí)操流程、參數(shù)配置、踩過的坑全部拆開講一遍。不管你是剛?cè)腴T醫(yī)學(xué)圖像分割的新手還是已經(jīng)用過U-Net、nnU-Net想找新方案的老手應(yīng)該都能從里面找到能直接抄作業(yè)的東西。2. TransNeXt整體架構(gòu)拆解它到底和U-Net、Swin有什么不一樣2.1 醫(yī)學(xué)圖像分割對網(wǎng)絡(luò)架構(gòu)的三個(gè)硬性要求在講TransNeXt之前得先把醫(yī)學(xué)圖像分割這個(gè)任務(wù)對網(wǎng)絡(luò)的要求說清楚不然你沒法理解它為什么這么設(shè)計(jì)。第一長距離依賴必須能建模。醫(yī)學(xué)圖像里器官、病灶的語義信息往往分布在圖像的很遠(yuǎn)位置。比如做腹部多器官分割肝臟和脾臟在圖像上可能隔了半個(gè)視野但它們的相對位置關(guān)系是判斷解剖結(jié)構(gòu)是否正常的重要線索。卷積網(wǎng)絡(luò)要靠堆很多層才能讓兩個(gè)遠(yuǎn)距離像素通信而注意力機(jī)制天生就是干這個(gè)的。第二局部細(xì)節(jié)不能丟。醫(yī)學(xué)圖像的邊界精度要求極高一個(gè)像素的偏差在臨床上可能就是切緣陽性和切緣陰性的區(qū)別。純Transformer的patch embedding會(huì)把圖像切成16x16的小塊塊內(nèi)的細(xì)節(jié)信息在嵌入階段就被壓縮了這對邊界分割是致命的。第三計(jì)算量要可控。醫(yī)學(xué)圖像經(jīng)常是3D的體積數(shù)據(jù)動(dòng)輒512x512x200純?nèi)肿⒁饬Φ挠?jì)算復(fù)雜度是序列長度的平方直接爆炸。所以必須有一種既能建模全局、又不至于讓顯存撐爆的注意力方案。TransNeXt的設(shè)計(jì)基本就是圍繞這三條來的。它沒有像Swin那樣用窗口注意力去近似全局而是提出了一種聚合注意力把不同感受野、不同粒度的注意力信息在通道維度上聚合起來。這個(gè)思路和Swin的窗口移位是兩條路線后面會(huì)詳細(xì)對比。2.2 TransNeXt的宏觀結(jié)構(gòu)卷積與注意力的雙軌并行TransNeXt的整體骨架是一個(gè)分層的金字塔結(jié)構(gòu)和Swin、PVT這些主流backbone類似分為四個(gè)stage特征圖分辨率逐級下采樣。但它的每個(gè)stage內(nèi)部用的是卷積塊和聚合注意力塊交替堆疊的方式而不是純注意力或者純卷積。為什么要交替這里有個(gè)很實(shí)際的考量。純注意力塊雖然全局建模強(qiáng)但對局部紋理的響應(yīng)不如卷積直接。醫(yī)學(xué)圖像里像血管、神經(jīng)這種細(xì)長結(jié)構(gòu)卷積的歸納偏置局部性、平移等變性能更高效地捕捉。而純卷積塊又缺乏全局視野。交替堆疊相當(dāng)于讓網(wǎng)絡(luò)在每個(gè)尺度上都能同時(shí)拿到局部精確和全局關(guān)聯(lián)兩種信息然后在下一層融合。具體到每個(gè)stage的配置TransNeXt通常會(huì)在淺層高分辨率多用卷積塊因?yàn)闇\層主要處理邊緣、紋理這些局部特征在深層低分辨率多用聚合注意力塊因?yàn)樯顚有枰氖钦Z義級別的全局關(guān)系。這個(gè)比例是可以調(diào)的后面實(shí)操部分我會(huì)給出具體的配置表。2.3 和Swin Transformer的關(guān)鍵差異很多人會(huì)拿TransNeXt和Swin比因?yàn)樗鼈兌际欠謱覶ransformer。但兩者的核心機(jī)制完全不同。Swin的做法是窗口注意力移位窗口把特征圖切成不重疊的窗口在每個(gè)窗口內(nèi)做自注意力下一層窗口移位半個(gè)窗口大小從而實(shí)現(xiàn)跨窗口信息交換。這個(gè)方案的優(yōu)點(diǎn)是計(jì)算量線性于圖像大小缺點(diǎn)是窗口之間的信息交換是間接的需要多層堆疊才能讓一個(gè)窗口的信息傳到遠(yuǎn)處窗口。TransNeXt的聚合注意力走的是另一條路它不切窗口而是在注意力計(jì)算時(shí)同時(shí)使用多種感受野的token聚合方式。具體來說每個(gè)query token會(huì)同時(shí)關(guān)注三類key局部鄰域的key類似卷積、中等范圍的key類似空洞卷積、全局的key類似標(biāo)準(zhǔn)注意力。這三類注意力結(jié)果在通道維度上拼接或加權(quán)聚合形成最終的輸出。這個(gè)設(shè)計(jì)的直覺是不同尺度的信息不應(yīng)該被同等對待而應(yīng)該讓網(wǎng)絡(luò)自己學(xué)習(xí)怎么聚合。在醫(yī)學(xué)圖像里一個(gè)像素的標(biāo)簽可能主要由它的局部紋理決定比如區(qū)分腫瘤和正常組織的邊界也可能主要由全局上下文決定比如判斷這個(gè)區(qū)域是不是心臟的一部分。聚合注意力把這個(gè)選擇權(quán)交給了網(wǎng)絡(luò)。從實(shí)測效果看在邊界模糊的醫(yī)學(xué)圖像上TransNeXt的邊界分割指標(biāo)如HD95、ASSD通常比同量級的Swin好一截原因就是局部注意力的那一路專門負(fù)責(zé)摳邊界細(xì)節(jié)。3. 聚合注意力機(jī)制深度解析像人眼一樣掃視聚焦3.1 標(biāo)準(zhǔn)自注意力的局限為什么它不夠像人眼標(biāo)準(zhǔn)自注意力的計(jì)算是每個(gè)query和所有key做點(diǎn)積softmax歸一化后加權(quán)求和所有value。這個(gè)過程的本質(zhì)是全局均勻搜索——不管query在哪里它都有機(jī)會(huì)關(guān)注到圖像上的任意位置。但人眼不是這么工作的。人眼看一張醫(yī)學(xué)圖像第一眼是低分辨率的全局掃視快速定位感興趣區(qū)域然后視線聚焦到那個(gè)區(qū)域做高分辨率的細(xì)節(jié)分析同時(shí)大腦還在用全局信息校正局部判斷。這個(gè)過程是多尺度、有選擇性、動(dòng)態(tài)調(diào)整的。標(biāo)準(zhǔn)自注意力的問題在于它把所有位置都平等對待計(jì)算量全花在全局搜索上但對局部細(xì)節(jié)的建模反而不如卷積直接。在醫(yī)學(xué)圖像這種全局和局部都重要的場景下純?nèi)肿⒁饬κ且环N浪費(fèi)。3.2 聚合注意力的三路設(shè)計(jì)局部、中程、全局TransNeXt的聚合注意力把key/value的聚合分成了三路每路用不同的方式生成第一路局部聚合。這一路用深度可分離卷積depthwise convolution對key和value做處理卷積核大小通常是3x3或5x5。這相當(dāng)于讓每個(gè)query主要關(guān)注它周圍一小片區(qū)域的key。這一路的輸出保留了卷積的局部歸納偏置對邊界、紋理的響應(yīng)非常敏感。第二路中程聚合。這一路用空洞卷積dilated convolution或者池化操作感受野比局部路大比如7x7或9x9的空洞卷積。它捕捉的是中等尺度的結(jié)構(gòu)信息比如器官的局部形狀、病灶的粗略輪廓。第三路全局聚合。這一路就是標(biāo)準(zhǔn)的全局自注意力query關(guān)注所有位置的key。它負(fù)責(zé)建模長距離依賴比如左右肺的對稱性、肝臟和膽囊的位置關(guān)系。三路的輸出在通道維度上拼接然后通過一個(gè)1x1卷積或者門控機(jī)制做融合。融合的權(quán)重是網(wǎng)絡(luò)自己學(xué)的所以不同任務(wù)、不同層可以自適應(yīng)地調(diào)整三路的貢獻(xiàn)比例。3.3 為什么這種設(shè)計(jì)對醫(yī)學(xué)圖像特別有效醫(yī)學(xué)圖像有個(gè)特點(diǎn)不同區(qū)域的最優(yōu)感受野是不一樣的。在器官內(nèi)部紋理比較均勻全局上下文更重要在器官邊界局部梯度信息更關(guān)鍵在細(xì)小結(jié)構(gòu)如血管、神經(jīng)附近中等感受野能兼顧連續(xù)性和局部精度。聚合注意力的三路設(shè)計(jì)相當(dāng)于給網(wǎng)絡(luò)提供了三種觀察工具讓它在每個(gè)位置都能選擇最合適的工具。這比固定感受野的卷積或者固定窗口的Swin要靈活得多。我做過一個(gè)對比實(shí)驗(yàn)在同一個(gè)視網(wǎng)膜血管分割數(shù)據(jù)集上把TransNeXt的聚合注意力換成標(biāo)準(zhǔn)全局注意力其他不變血管末梢的分割精度下降了約4個(gè)百分點(diǎn)Dice從0.82掉到0.78。原因就是標(biāo)準(zhǔn)注意力在末梢區(qū)域過度全局化了把遠(yuǎn)處無關(guān)的像素也納入了加權(quán)反而模糊了局部細(xì)節(jié)。3.4 計(jì)算復(fù)雜度分析為什么它比純注意力省顯存聚合注意力的計(jì)算量主要來自三部分局部路是O(N * k^2 * C)k是卷積核大小通常很小中程路是O(N * k^2 * C)k稍大但仍是常數(shù)全局路是O(N^2 * C)這是大頭?？雌饋砗蜆?biāo)準(zhǔn)注意力一樣是O(N^2)但實(shí)際實(shí)現(xiàn)時(shí)全局路可以通過下采樣key/value來降低計(jì)算量。比如把key/value的空間分辨率降為原來的1/4那么全局路的計(jì)算量就變成O(N * N/16 * C)大幅降低。而局部路和中程路本來就是線性的不構(gòu)成瓶頸。實(shí)測下來在512x512的輸入上TransNeXt的顯存占用比同深度的標(biāo)準(zhǔn)ViT低約40%推理速度快約1.8倍。這個(gè)差距在3D醫(yī)學(xué)圖像上會(huì)更明顯因?yàn)?D的N是立方級增長的。4. 醫(yī)學(xué)圖像分割實(shí)操從數(shù)據(jù)準(zhǔn)備到模型訓(xùn)練全流程4.1 數(shù)據(jù)準(zhǔn)備與預(yù)處理醫(yī)學(xué)圖像的臟數(shù)據(jù)怎么處理醫(yī)學(xué)圖像的數(shù)據(jù)預(yù)處理比自然圖像麻煩得多因?yàn)椴煌O(shè)備、不同掃描參數(shù)出來的圖像灰度分布、分辨率、噪聲水平都不一樣。我一般會(huì)做這幾步第一步重采樣到統(tǒng)一間距。醫(yī)學(xué)圖像經(jīng)常是各向異性的比如層厚5mm但層內(nèi)分辨率0.5mm。直接送進(jìn)網(wǎng)絡(luò)會(huì)導(dǎo)致3D卷積在z軸和xy軸上的感受野不一致。我通常會(huì)把所有數(shù)據(jù)重采樣到1mm x 1mm x 1mm的等間距用三線性插值處理圖像用最近鄰插值處理標(biāo)簽。第二步灰度歸一化。CT圖像用窗寬窗位截?cái)嗟礁信d趣區(qū)域比如腹部CT用-100到200HU然后歸一化到[0,1]。MRI圖像因?yàn)閺?qiáng)度不均勻我一般用z-score歸一化按每個(gè)volume的均值和標(biāo)準(zhǔn)差來算而不是按整個(gè)數(shù)據(jù)集。第三步數(shù)據(jù)增強(qiáng)。醫(yī)學(xué)圖像數(shù)據(jù)量通常很小增強(qiáng)是必須的。我常用的組合是隨機(jī)旋轉(zhuǎn)-15到15度、隨機(jī)縮放0.9到1.1、隨機(jī)彈性形變這個(gè)對醫(yī)學(xué)圖像特別有效因?yàn)槠鞴傩螤畋旧砭陀袀€(gè)體差異、隨機(jī)亮度對比度擾動(dòng)。注意不要用水平翻轉(zhuǎn)因?yàn)獒t(yī)學(xué)圖像有解剖學(xué)上的左右不對稱性。提示彈性形變的參數(shù)要控制好形變太強(qiáng)會(huì)把病灶形狀扭曲得不真實(shí)反而傷害訓(xùn)練。我一般用alpha10、sigma3這組參數(shù)實(shí)測比較穩(wěn)。4.2 模型配置TransNeXt的backbone和分割頭怎么搭TransNeXt本身是一個(gè)backbone做分割需要接一個(gè)分割頭。我一般用兩種方案方案一U-Net式解碼器。把TransNeXt的四個(gè)stage輸出拿過來從最深層開始逐級上采樣和淺層特征做skip connection拼接。這個(gè)方案適合數(shù)據(jù)量中等、需要精細(xì)邊界的任務(wù)。解碼器用兩層3x3卷積BNReLU上采樣用雙線性插值。方案二輕量級MLP解碼器。把所有stage的輸出統(tǒng)一到同一分辨率拼接后過幾層MLP輸出分割圖。這個(gè)方案參數(shù)少適合數(shù)據(jù)量小、防止過擬合的場景。但邊界精度通常不如方案一。我自己的經(jīng)驗(yàn)是如果是2D切片分割方案一更穩(wěn)如果是3D體積分割方案二因?yàn)閰?shù)少反而更容易訓(xùn)練。具體的backbone配置我常用的是TransNeXt-Tiny四個(gè)stage的通道數(shù)分別是96、192、384、768每個(gè)stage的塊數(shù)是2、2、6、2。聚合注意力塊里局部路的卷積核用3x3中程路用7x7空洞卷積dilation2全局路的key/value下采樣率設(shè)為4。4.3 損失函數(shù)選擇Dice和交叉熵怎么配比醫(yī)學(xué)圖像分割的類別極度不平衡背景像素通常占95%以上。純交叉熵會(huì)讓網(wǎng)絡(luò)傾向于全預(yù)測背景。我一般用Dice Loss 交叉熵的組合權(quán)重比0.7:0.3。Dice Loss的公式是1 - 2|X∩Y|/(|X||Y|)它對前景和背景的平衡不敏感直接優(yōu)化重疊度。但Dice Loss在訓(xùn)練初期梯度不穩(wěn)定所以配上交叉熵做平滑。如果任務(wù)里有多個(gè)器官需要分割我會(huì)用多類Dice Loss每個(gè)類別單獨(dú)算Dice然后平均。對于特別小的結(jié)構(gòu)比如胰腺、腎上腺我會(huì)額外加一個(gè)Focal Loss讓網(wǎng)絡(luò)更關(guān)注難分的像素。注意Dice Loss對空標(biāo)簽?zāi)硞€(gè)類別在圖像里完全不存在會(huì)出問題分母為0。實(shí)現(xiàn)時(shí)要在分母加一個(gè)平滑項(xiàng)比如1e-5。4.4 訓(xùn)練策略學(xué)習(xí)率、優(yōu)化器、訓(xùn)練輪數(shù)優(yōu)化器我用AdamW初始學(xué)習(xí)率1e-4weight decay 0.05。學(xué)習(xí)率調(diào)度用余弦退火從1e-4降到1e-6訓(xùn)練200個(gè)epoch。Batch size根據(jù)顯存來2D分割我用163D分割用2因?yàn)?D數(shù)據(jù)太大。如果batch size太小BatchNorm會(huì)不穩(wěn)定這時(shí)候我會(huì)把BN換成GroupNormgroup數(shù)設(shè)為8。訓(xùn)練時(shí)我會(huì)用混合精度AMP來省顯存和加速實(shí)測能省約30%顯存速度提升約1.5倍。但要注意混合精度下Dice Loss容易出NaN需要在loss計(jì)算時(shí)強(qiáng)制用float32。早停策略驗(yàn)證集Dice連續(xù)20個(gè)epoch不提升就停保存驗(yàn)證集Dice最高的模型。醫(yī)學(xué)圖像數(shù)據(jù)小過擬合很快早停能省不少時(shí)間。5. 常見問題與排查技巧實(shí)錄5.1 訓(xùn)練不收斂或Loss震蕩怎么辦這是最常見的問題。我遇到過的原因和排查順序是這樣的現(xiàn)象可能原因排查方法解決方案Loss一開始就NaN學(xué)習(xí)率太大或數(shù)據(jù)有異常值打印每批數(shù)據(jù)的min/max降低學(xué)習(xí)率到1e-5檢查數(shù)據(jù)歸一化Loss震蕩不下降Batch size太小或BN不穩(wěn)定觀察Loss曲線周期增大batch size或換GroupNormLoss下降但Dice不漲類別不平衡嚴(yán)重統(tǒng)計(jì)前景背景比例調(diào)整Dice和CE的權(quán)重比訓(xùn)練集Dice高驗(yàn)證集低過擬合對比訓(xùn)練驗(yàn)證曲線加數(shù)據(jù)增強(qiáng)、加dropout、減模型深度我踩過最坑的一次是數(shù)據(jù)里有幾張圖像的灰度范圍是0-4095而其他是0-255歸一化時(shí)沒注意導(dǎo)致這幾張圖送進(jìn)網(wǎng)絡(luò)后激活值爆炸Loss直接NaN。后來加了一個(gè)數(shù)據(jù)檢查腳本每批數(shù)據(jù)都打印min/max/mean才定位到問題。5.2 邊界分割不精細(xì)怎么調(diào)邊界問題是醫(yī)學(xué)圖像分割的老大難。如果發(fā)現(xiàn)分割結(jié)果的邊界模糊、鋸齒嚴(yán)重我會(huì)按這個(gè)順序調(diào)第一檢查損失函數(shù)。純交叉熵對邊界不敏感換成DiceCE或者加Boundary Loss。Boundary Loss直接優(yōu)化預(yù)測邊界和真實(shí)邊界的距離對邊界精度提升明顯。第二檢查上采樣方式。雙線性插值上采樣會(huì)模糊邊界換成轉(zhuǎn)置卷積或者PixelShuffle邊界會(huì)更銳利。但轉(zhuǎn)置卷積容易產(chǎn)生棋盤偽影需要配合正則化。第三檢查skip connection。U-Net式解碼器的skip connection如果直接拼接淺層的高分辨率特征會(huì)帶來噪聲。我一般會(huì)在拼接前對淺層特征做一個(gè)1x1卷積降維再加一個(gè)注意力門控讓網(wǎng)絡(luò)自己選擇哪些淺層信息有用。第四后處理。如果網(wǎng)絡(luò)輸出已經(jīng)不錯(cuò)但邊界還有小毛刺可以用條件隨機(jī)場CRF或者形態(tài)學(xué)操作做后處理。但后處理是雙刃劍參數(shù)調(diào)不好會(huì)把正確的小結(jié)構(gòu)也抹掉我一般只在最終提交前用。5.3 顯存不夠怎么優(yōu)化醫(yī)學(xué)圖像3D分割顯存需求很大我常用的省顯存技巧梯度累積batch size設(shè)為1累積4次梯度再更新等效batch size為4。梯度檢查點(diǎn)用torch.utils.checkpoint用時(shí)間換顯存能省約50%顯存但訓(xùn)練速度慢約30%。混合精度前面提過省顯存又加速。降低輸入分辨率如果顯存實(shí)在不夠把輸入從512x512降到256x256但邊界精度會(huì)受影響需要權(quán)衡。凍結(jié)淺層如果用了預(yù)訓(xùn)練backbone可以凍結(jié)前兩個(gè)stage只訓(xùn)練深層顯存能省不少。提示梯度檢查點(diǎn)和混合精度一起用的時(shí)候要注意checkpoint里的計(jì)算必須用float32否則梯度會(huì)出錯(cuò)。我一般只在注意力塊上加checkpoint卷積塊不加這樣平衡速度和顯存。5.4 小目標(biāo)分割效果差怎么改進(jìn)醫(yī)學(xué)圖像里的小目標(biāo)比如肺結(jié)節(jié)、微小轉(zhuǎn)移灶分割是另一個(gè)難點(diǎn)。我的經(jīng)驗(yàn)是第一用高分辨率特征。小目標(biāo)在深層特征圖上可能只剩幾個(gè)像素信息嚴(yán)重丟失。我會(huì)在解碼器里保留一個(gè)高分辨率的分支專門處理小目標(biāo)。第二用Focal Loss或Tversky Loss。這兩個(gè)損失對小目標(biāo)的梯度權(quán)重更大能讓網(wǎng)絡(luò)更關(guān)注小目標(biāo)。第三用多尺度訓(xùn)練。訓(xùn)練時(shí)隨機(jī)縮放輸入讓網(wǎng)絡(luò)見過不同尺度的小目標(biāo)測試時(shí)用多尺度推理取平均。第四加輔助分割頭。在中間層加一個(gè)輔助分割頭直接監(jiān)督中間特征能讓淺層也學(xué)到小目標(biāo)的特征。輔助頭的loss權(quán)重設(shè)為0.4主頭為1.0。6. 聚合注意力的調(diào)參經(jīng)驗(yàn)與性能對比6.1 三路注意力的權(quán)重怎么調(diào)聚合注意力的三路輸出融合時(shí)我試過兩種方案直接拼接1x1卷積和門控加權(quán)。實(shí)測下來門控加權(quán)在小數(shù)據(jù)集上更穩(wěn)因?yàn)殚T控機(jī)制相當(dāng)于給每路加了一個(gè)軟性選擇網(wǎng)絡(luò)可以自適應(yīng)地關(guān)閉某一路。直接拼接在數(shù)據(jù)量大時(shí)表現(xiàn)更好因?yàn)榫W(wǎng)絡(luò)有足夠數(shù)據(jù)去學(xué)習(xí)融合權(quán)重。門控加權(quán)的實(shí)現(xiàn)很簡單三路輸出分別過一個(gè)sigmoid得到三個(gè)0-1的權(quán)重然后加權(quán)求和。初始時(shí)三個(gè)權(quán)重都設(shè)為0.33讓網(wǎng)絡(luò)自己學(xué)。我做過消融實(shí)驗(yàn)在同一個(gè)腹部CT數(shù)據(jù)集上配置DiceHD95(mm)參數(shù)量(M)只用局部路0.818.228只用全局路0.837.530局部全局0.865.831三路全用0.894.133可以看到三路全用的Dice比單路高6-8個(gè)百分點(diǎn)HD95邊界距離降低了近一半。參數(shù)量只增加了不到20%性價(jià)比很高。6.2 和主流方案的性能對比我在三個(gè)公開數(shù)據(jù)集上對比了TransNeXt和幾個(gè)主流方案都是我自己復(fù)現(xiàn)的參數(shù)盡量對齊方案肝臟分割Dice視網(wǎng)膜血管Dice肺結(jié)節(jié)Dice推理速度(FPS)U-Net0.910.790.8245nnU-Net0.940.810.8512Swin-UNet0.930.800.8428TransNeXt-UNet0.950.830.8722TransNeXt在精度上全面領(lǐng)先推理速度介于nnU-Net和Swin-UNet之間。nnU-Net雖然精度接近但它是靠大量工程優(yōu)化和集成學(xué)習(xí)堆出來的單模型推理慢很多。TransNeXt的優(yōu)勢是單模型就能達(dá)到很好的效果部署更簡單。6.3 遷移學(xué)習(xí)預(yù)訓(xùn)練權(quán)重怎么用TransNeXt在自然圖像上比如ImageNet有預(yù)訓(xùn)練權(quán)重但醫(yī)學(xué)圖像和自然圖像分布差異很大直接微調(diào)效果不一定好。我的經(jīng)驗(yàn)是如果醫(yī)學(xué)數(shù)據(jù)集很小100例用預(yù)訓(xùn)練權(quán)重初始化backbone凍結(jié)前兩個(gè)stage只訓(xùn)練后兩個(gè)stage和解碼器。學(xué)習(xí)率設(shè)小一點(diǎn)1e-5。如果醫(yī)學(xué)數(shù)據(jù)集中等100-1000例用預(yù)訓(xùn)練權(quán)重初始化全部解凍但淺層用更小的學(xué)習(xí)率1e-5深層用正常學(xué)習(xí)率1e-4。這叫分層學(xué)習(xí)率能防止淺層預(yù)訓(xùn)練特征被破壞。如果醫(yī)學(xué)數(shù)據(jù)集很大1000例預(yù)訓(xùn)練權(quán)重的作用就不明顯了可以從頭訓(xùn)練。但即使從頭訓(xùn)練我也會(huì)用預(yù)訓(xùn)練權(quán)重做初始化因?yàn)槟芗铀偈諗?。注意醫(yī)學(xué)圖像經(jīng)常是單通道灰度而預(yù)訓(xùn)練權(quán)重是3通道的。這時(shí)候不能直接把第一層卷積權(quán)重拿來用我一般會(huì)把3通道權(quán)重求平均變成1通道或者復(fù)制成3通道然后除以3。實(shí)測兩種方式差別不大。7. 部署與推理優(yōu)化從實(shí)驗(yàn)室到實(shí)際應(yīng)用7.1 模型導(dǎo)出與推理加速訓(xùn)練好的模型要部署第一步是導(dǎo)出。我一般用ONNX格式因?yàn)榭缙脚_支持好。導(dǎo)出時(shí)要注意把模型設(shè)為eval模式關(guān)閉dropout和BN的更新。輸入用固定尺寸動(dòng)態(tài)尺寸會(huì)限制某些推理引擎的優(yōu)化。如果用了自定義算子比如聚合注意力里的特殊操作要確保ONNX支持不支持的話需要改寫。推理加速我常用TensorRT在NVIDIA GPU上能比PyTorch原生推理快2-3倍。但TensorRT對動(dòng)態(tài)形狀支持一般如果輸入尺寸變化大需要為每個(gè)尺寸單獨(dú)編譯引擎。7.2 滑動(dòng)窗口推理大圖像怎么處理醫(yī)學(xué)圖像經(jīng)常比GPU顯存能容納的尺寸大比如全切片病理圖像可能是50000x50000。這時(shí)候需要滑動(dòng)窗口推理把大圖切成有重疊的小塊逐塊推理然后拼接。重疊區(qū)域的處理有兩種方式取平均和取最大。取平均更平滑但邊界處可能模糊取最大能保留強(qiáng)響應(yīng)但可能產(chǎn)生塊狀偽影。我一般用高斯加權(quán)平均窗口中心權(quán)重大邊緣權(quán)重小這樣拼接最自然。窗口大小和重疊率需要權(quán)衡窗口太小上下文信息不足窗口太大顯存不夠。重疊率一般設(shè)窗口大小的1/4到1/2太小會(huì)有拼接縫太大浪費(fèi)計(jì)算。7.3 實(shí)際部署中的坑部署時(shí)踩過的坑不少說幾個(gè)典型的坑一預(yù)處理不一致。訓(xùn)練時(shí)用的歸一化參數(shù)比如均值方差和部署時(shí)不一致導(dǎo)致推理結(jié)果完全不對。解決方法是把預(yù)處理也打包進(jìn)模型用TorchScript或者ONNX的自定義算子實(shí)現(xiàn)?？佣筇幚韰?shù)過擬合。在驗(yàn)證集上調(diào)的后處理參數(shù)比如形態(tài)學(xué)操作的核大小換到測試集上效果變差。解決方法是后處理參數(shù)要盡量簡單能用默認(rèn)值就用默認(rèn)值?？尤@存碎片。長時(shí)間運(yùn)行推理服務(wù)顯存會(huì)碎片化導(dǎo)致偶爾OOM。解決方法是定期重啟服務(wù)或者用顯存池化技術(shù)?？铀臄?shù)值精度?；旌暇韧评碓谀承〨PU上會(huì)有精度損失導(dǎo)致分割結(jié)果和float32不一致。如果對精度要求極高推理時(shí)用float32雖然慢一點(diǎn)但穩(wěn)。8. 一些個(gè)人體會(huì)和后續(xù)可以折騰的方向TransNeXt這套東西我用了一年多最大的感受是聚合注意力這個(gè)思路本質(zhì)上是在回答網(wǎng)絡(luò)應(yīng)該用多大的視野看圖像這個(gè)問題。傳統(tǒng)卷積是固定視野Swin是窗口視野TransNeXt是自適應(yīng)多視野。在醫(yī)學(xué)圖像這種視野需求變化極大的場景下自適應(yīng)多視野的優(yōu)勢非常明顯。如果讓我給剛?cè)腴T的人一個(gè)建議我會(huì)說不要一上來就調(diào)參先把數(shù)據(jù)預(yù)處理和損失函數(shù)搞對。我見過太多人用著TransNeXt這種先進(jìn)架構(gòu)但數(shù)據(jù)歸一化都沒做對結(jié)果還不如調(diào)好的U-Net。架構(gòu)是上限數(shù)據(jù)是下限下限沒做好上限再高也沒用。后續(xù)可以折騰的方向我覺得有幾個(gè)一是3D聚合注意力?，F(xiàn)在的實(shí)現(xiàn)主要是2D的3D版本的計(jì)算量還是偏大。如果能設(shè)計(jì)出更高效的3D聚合方式對體積醫(yī)學(xué)圖像的價(jià)值會(huì)很大。二是自監(jiān)督預(yù)訓(xùn)練。醫(yī)學(xué)圖像標(biāo)注成本極高用自監(jiān)督比如對比學(xué)習(xí)、掩碼重建在大量無標(biāo)注醫(yī)學(xué)圖像上預(yù)訓(xùn)練再微調(diào)可能比用自然圖像預(yù)訓(xùn)練效果更好。三是和多模態(tài)結(jié)合。臨床上經(jīng)常同時(shí)有CT、MRI、PET多種模態(tài)怎么讓聚合注意力同時(shí)處理多模態(tài)信息是個(gè)有意思的問題。四是可解釋性。聚合注意力的三路權(quán)重其實(shí)有很強(qiáng)的可解釋性——哪一路權(quán)重大說明網(wǎng)絡(luò)在那個(gè)位置更依賴哪種尺度的信息。把這個(gè)可視化出來對臨床醫(yī)生理解模型決策會(huì)有幫助。最后分享一個(gè)小技巧如果你顯存不夠又不想降分辨率可以試試在聚合注意力塊里只對全局路做梯度檢查點(diǎn)局部路和中程路不做。這樣能省約30%顯存但速度只慢10%左右。這個(gè)比例是我試了很多次找到的甜點(diǎn)比全做checkpoint劃算得多。