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