簡(jiǎn)介面向Python開發(fā)者的圖像分割實(shí)戰(zhàn)資料圍繞U-Net模型講解如何完成像素級(jí)分割任務(wù)適用于遙感影像、醫(yī)學(xué)圖像等場(chǎng)景。壓縮包共21個(gè)文件、約5.6MB內(nèi)附1個(gè)Python腳本平滑切片預(yù)測(cè)、多張jpg/png樣例圖、2個(gè)gif動(dòng)態(tài)對(duì)比圖以及Markdown說明文檔便于對(duì)照代碼與效果圖理解模型推理過程。已有10157人學(xué)習(xí)下載。資料內(nèi)容涵蓋U-Net收縮/擴(kuò)展路徑原理、數(shù)據(jù)掩模準(zhǔn)備、交叉熵與Dice損失選擇、優(yōu)化器配置等關(guān)鍵知識(shí)點(diǎn)代碼庫(kù)Smoothly-Blend-Image-Patches-master實(shí)現(xiàn)了圖像分塊平滑融合可用于消除切片邊界偽影提升分割效果。通過閱讀README并運(yùn)行腳本讀者能掌握從模型構(gòu)建、訓(xùn)練到預(yù)測(cè)的完整流程還可借助樣例圖像直觀了解U-Net在衛(wèi)星圖像分割中的實(shí)際表現(xiàn)適合希望以Python快速落地圖像分割項(xiàng)目的中級(jí)開發(fā)者。1. UNet 圖像分割把像素分類這件事講透圖像分割和分類最大的區(qū)別在于分類只回答“這張圖里有什么”而分割要回答“這個(gè)物體在哪、邊界在哪里、和旁邊的物體怎么分開”。UNet 之所以成為分割任務(wù)里的常青樹是因?yàn)樗诰幋a器-解碼器結(jié)構(gòu)上加了跳躍連接用很小的參數(shù)量同時(shí)保住了全局語(yǔ)義和局部細(xì)節(jié)——我做過不少分割項(xiàng)目從路面裂縫到廣告牌檢測(cè)最后落地的方案里 UNet 和它的變體占了大多數(shù)。這個(gè)資源提供了一個(gè)可以直接跑的 Python 實(shí)現(xiàn)適合剛接觸分割的入門者也適合需要快速做可行性驗(yàn)證的從業(yè)者。它能幫你把一個(gè)圖像分割模型從數(shù)據(jù)準(zhǔn)備、訓(xùn)練到預(yù)測(cè)完整跑通整個(gè)鏈路不超過 200 行核心代碼。2. UNet 網(wǎng)絡(luò)結(jié)構(gòu)為什么所有分割任務(wù)都繞不開它2.1 編碼器-解碼器骨架和跳躍連接UNet 的結(jié)構(gòu)可以拆成三段看編碼器、解碼器和跳躍連接。編碼器就是不斷卷積加下采樣把輸入從 512×512 壓到 16×16這個(gè)過程中通道數(shù)從 3 漲到 512特征圖越來越抽象對(duì)應(yīng)的是“語(yǔ)義信息”。解碼器則反向操作一步步上采樣把特征圖恢復(fù)回原分辨率這時(shí)候網(wǎng)絡(luò)要學(xué)會(huì)把抽象語(yǔ)義和空間位置對(duì)應(yīng)起來。如果沒有跳躍連接解碼器在恢復(fù)細(xì)節(jié)時(shí)會(huì)嚴(yán)重丟信息因?yàn)橄虏蓸影芽臻g位置徹底打散了。跳躍連接做的事情非常樸素把編碼器某一層的特征圖直接拼到解碼器對(duì)應(yīng)層上。例如第 3 層下采樣前的輸出是 64×64×256解碼器在同一尺度上恢復(fù)時(shí)就把這兩個(gè)特征圖在通道維度上拼起來得到 64×64×512然后再卷積。這讓解碼器既能看到高層語(yǔ)義又能直接訪問淺層的邊緣紋理信息。實(shí)際訓(xùn)練中的感受是去掉跳躍連接后模型分割出來的物體邊緣會(huì)出現(xiàn)明顯的“膨脹”或“鋸齒”加上之后效果會(huì)立刻干凈很多。這個(gè)設(shè)計(jì)讓 UNet 在小數(shù)據(jù)集上表現(xiàn)特別好因?yàn)樗恍枰獜牧銓W(xué)習(xí)“邊緣在哪里”這件事情圖像的低層結(jié)構(gòu)已經(jīng)通過跳躍連接直接送到解碼器了。2.2 深度、通道數(shù)與輸入尺寸的選型UNet 的原始論文用的是 32 作為初始通道數(shù)每下采樣一次通道翻倍直到 512。這個(gè)配置并不是拍腦袋定的它的出發(fā)點(diǎn)是讓每一層都有足夠的表達(dá)能力去編碼對(duì)應(yīng)尺度的信息。實(shí)際使用時(shí)要看“你的物體有多大”這個(gè)關(guān)鍵問題。如果做的是車道線分割或裂縫檢測(cè)物體細(xì)長(zhǎng)輸入尺寸建議至少 512×512初始通道數(shù)可以降到 16因?yàn)?32 通道在推理時(shí)會(huì)明顯拖慢速度而精度提升很有限。如果做的是地毯、廣告牌這類大塊面物體256×256 輸入加初始通道 16 已經(jīng)足夠硬上 512 輸入只會(huì)增加顯存占用。編碼器深度的選擇和一些常見的誤區(qū)有關(guān)。很多人以為層數(shù)越多效果越好但 UNet 下采樣到 32×32 以下時(shí)對(duì)大多數(shù)分割任務(wù)已經(jīng)過于抽象反而會(huì)讓小目標(biāo)直接消失。我在工程里常用的判斷標(biāo)準(zhǔn)是感受野至少要覆蓋目標(biāo)最大尺寸的一半。用 512×512 輸入、5 次下采樣的 UNet最后一層的感受野覆蓋整個(gè)圖但對(duì) 10×10 像素的小目標(biāo)來說它在下采樣第三次時(shí)就已經(jīng)被壓縮成了 1-2 個(gè)像素——所以遇到底層細(xì)節(jié)任務(wù)時(shí)我會(huì)減少下采樣次數(shù)或者跳過最后一個(gè)跳躍連接讓網(wǎng)絡(luò)在保留空間信息的前提下做特征融合比盲目加深有效得多。2.3 從 ResNet 編碼器到輕量化變體的改進(jìn)邏輯標(biāo)準(zhǔn) UNet 的編碼器是重復(fù)的卷積塊它的表達(dá)能力和 ResNet34 這類預(yù)訓(xùn)練骨干相比存在明顯差距。使用 ImageNet 預(yù)訓(xùn)練的 ResNet34 作為骨干遷移學(xué)習(xí)的效果幾乎總是優(yōu)于隨機(jī)初始化的標(biāo)準(zhǔn) UNet尤其當(dāng)訓(xùn)練數(shù)據(jù)只有幾百?gòu)垐D時(shí)。原因很直觀預(yù)訓(xùn)練骨干已經(jīng)學(xué)到了通用的邊緣、紋理、顏色特征分割需要的底層視覺能力不需要從零學(xué)。代碼里把 UNet 的編碼器部分替換成 ResNet34解碼器保持原樣然后把跳躍連接的通道數(shù)改成 ResNet 各層的輸出通道即可。輕量化方向則是另一個(gè)極端對(duì)實(shí)時(shí)性要求高的場(chǎng)景MobilenetV2 作為編碼器的 UNet 是常見選擇。它用深度可分離卷積把計(jì)算量大幅降低。實(shí)際使用時(shí)要注意MobilenetV2 的特征圖通道排布和標(biāo)準(zhǔn) UNet 不同跳躍連接拼接時(shí)要讀一下各層的 out_channels而不是直接照搬源碼。至于注意力機(jī)制像在跳躍連接后面加一個(gè)簡(jiǎn)單的 SE 模塊或 CBAM確實(shí)能提升邊界處的精度但代價(jià)是參數(shù)和推理耗時(shí)增加在工程里我通常只在最后一層跳躍連接上加而不是每層都加。3. 數(shù)據(jù)準(zhǔn)備與預(yù)處理決定效果的七成因素3.1 數(shù)據(jù)目錄組織與標(biāo)簽格式分割數(shù)據(jù)集的標(biāo)配是“圖像 掩膜”的組合。常見做法是按images和masks兩個(gè)目錄平鋪存放文件名一一對(duì)應(yīng)。掩膜有兩種格式一種是 PNG 索引圖像素值為 0、1、2 這樣的類別 ID另一種是灰度圖或 RGB 圖用 0 和 255 表示前景背景。這看起來是小事但數(shù)據(jù)加載代碼必須知道自己處理的是哪一種否則后面訓(xùn)練會(huì)在標(biāo)簽讀取時(shí)吃大虧。我的經(jīng)驗(yàn)是先寫一個(gè)數(shù)據(jù)檢查函數(shù)把每張掩膜的unique值打出來看看。索引圖應(yīng)該是[0, 1, 2]這種連續(xù)小整數(shù)而掩膜圖可能是[0, 255]這時(shí)需要在加載時(shí)做一個(gè)mask // 255的操作把 255 映射回 1。還有一個(gè)常見的坑是某些標(biāo)注工具導(dǎo)出的掩膜帶有調(diào)色板模式用cv2.imread讀出來會(huì)是三通道。這時(shí)候要cv2.imread(path, cv2.IMREAD_GRAYSCALE)強(qiáng)制讀成灰度圖否則三通道的掩膜會(huì)被誤當(dāng)作 RGB 圖像參與訓(xùn)練損失函數(shù)計(jì)算時(shí)直接出錯(cuò)。建議第一次跑通前先打印三張樣本的 shape 和 dtype 確認(rèn)。數(shù)據(jù)增強(qiáng)的順序也常有人搞反先對(duì)圖像做幾何變換再對(duì)掩膜做完全相同的變換最后再分別做歸一化。這個(gè)過程中最怕的是用torchvision.transforms里的隨機(jī)類變換它對(duì)圖像和掩膜是分別隨機(jī)化的導(dǎo)致圖像旋轉(zhuǎn)了 10 度而掩膜旋轉(zhuǎn)了 15 度。我的做法是自定義一個(gè)SegPairTransform類在同一個(gè)seed下生成隨機(jī)參數(shù)然后分別應(yīng)用到圖像和掩膜上保證它們空間變換完全一致。做分割項(xiàng)目時(shí)數(shù)據(jù)增強(qiáng)的 bug 往往是訓(xùn)練時(shí) loss 正常下降、驗(yàn)證時(shí)指標(biāo)卻一直上不去的頭號(hào)元兇因?yàn)樗褬?biāo)簽錯(cuò)位了卻不報(bào)錯(cuò)。3.2 加載器與增強(qiáng)策略的完整代碼以下是我在分割項(xiàng)目里常用的數(shù)據(jù)加載器實(shí)現(xiàn)按“讀圖 → 統(tǒng)一尺寸 → 成對(duì)增強(qiáng) → 張量化”的流程走import cv2 import numpy as np import torch from torch.utils.data import Dataset from torchvision import transforms class SegDataset(Dataset): def __init__(self, img_dir, mask_dir, size(512, 512), augmentFalse): self.img_paths sorted(list(img_dir.glob(*.png))) # 按文件名排序 self.mask_paths sorted(list(mask_dir.glob(*.png))) self.augment augment self.resize_h, self.resize_w size assert len(self.img_paths) len(self.mask_paths), 圖像和掩膜數(shù)量不一致 def __getitem__(self, idx): img cv2.imread(str(self.img_paths[idx])) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) mask cv2.imread(str(self.mask_paths[idx]), cv2.IMREAD_GRAYSCALE) mask mask // 255 # 255 映射到 1 img cv2.resize(img, (self.resize_w, self.resize_h), interpolationcv2.INTER_LINEAR) mask cv2.resize(mask, (self.resize_w, self.resize_h), interpolationcv2.INTER_NEAREST) if self.augment: img, mask self._pair_augment(img, mask) img torch.from_numpy(img).permute(2, 0, 1).float() / 255.0 mask torch.from_numpy(mask).long() return img, mask def _pair_augment(self, img, mask): # 同一個(gè)隨機(jī)種子保證空間變換一致 seed np.random.randint(0, 2**31) cv2.setRNGSeed(seed) if np.random.random() 0.5: img cv2.flip(img, 1) mask cv2.flip(mask, 1) cv2.setRNGSeed(seed) # 重置種子讓旋轉(zhuǎn)參數(shù)保持一致 angle np.random.uniform(-15, 15) h, w img.shape[:2] M cv2.getRotationMatrix2D((w / 2, h / 2), angle, 1.0) img cv2.warpAffine(img, M, (w, h), flagscv2.INTER_LINEAR) mask cv2.warpAffine(mask, M, (w, h), flagscv2.INTER_NEAREST) return img, mask這段代碼里最關(guān)鍵的是mask // 255和重采樣方式。INTER_NEAREST是掩膜縮放時(shí)唯一推薦的方式因?yàn)樗粫?huì)產(chǎn)生插值導(dǎo)致的偽類別比如背景經(jīng)過INTER_LINEAR后出現(xiàn) 128 這樣不屬于任何類別的中間值。圖像則用INTER_LINEAR保持邊緣平滑。旋轉(zhuǎn)使用的cv2.setRNGSeed是為了讓圖像和掩膜用同一組隨機(jī)參數(shù)否則分割標(biāo)簽會(huì)在空間上錯(cuò)位。3.3 數(shù)據(jù)量不夠時(shí)怎么補(bǔ)救標(biāo)注數(shù)據(jù)不夠是分割項(xiàng)目的常態(tài)幾百?gòu)垐D也同樣能訓(xùn)練出一個(gè)能用的模型。首要手段是離線增強(qiáng)除了在線增強(qiáng)的隨機(jī)翻轉(zhuǎn)旋轉(zhuǎn)再用 HSV 空間的隨機(jī)色度抖動(dòng)增加顏色多樣性。另一個(gè)被驗(yàn)證有效的做法是使用預(yù)訓(xùn)練編碼器把 UNet 的骨干換成 ResNet34 并用 ImageNet 權(quán)重初始化訓(xùn)練只需調(diào)解碼器部分幾百?gòu)垐D完全足夠。還有一個(gè)方向是偽標(biāo)簽先用當(dāng)前模型對(duì)無標(biāo)簽數(shù)據(jù)做預(yù)測(cè)挑置信度高的區(qū)域作為訓(xùn)練數(shù)據(jù)——這個(gè)辦法在沒有標(biāo)注標(biāo)注人力時(shí)的確有效也直接幫助后續(xù)預(yù)測(cè)階段積累樣本。如果連幾百?gòu)垐D都沒有就先別談?dòng)?xùn)練直接把任務(wù)降級(jí)為目標(biāo)檢測(cè)或者用現(xiàn)成的分割模型先做一輪預(yù)標(biāo)注再人工修正效率遠(yuǎn)高于手工從頭標(biāo)注。4. 訓(xùn)練配置與計(jì)算損失函數(shù)、學(xué)習(xí)率和顯存管理4.1 損失函數(shù)的選擇邏輯UNet 最常用的損失函數(shù)組合是CrossEntropyLoss和DiceLoss兩類任務(wù)有明確的適用邊界。如果目標(biāo)是背景占比大而前景占比很小比如裂縫、血管CrossEntropyLoss會(huì)直接偏向背景類網(wǎng)絡(luò)預(yù)測(cè)結(jié)果可能全圖都是背景。DiceLoss 直接優(yōu)化前景區(qū)域的覆蓋率在這種類別極端不均衡的場(chǎng)景下效果更穩(wěn)定。如果目標(biāo)是前景和背景比例接近的比如地毯、廣告牌這種大塊物體BCEWithLogitsLoss就夠用配一個(gè) 0.2 到 0.4 的類別權(quán)重即可。工程里我一般不用單一的 DiceLoss而是用 BCE 和 Dice 加權(quán)相加經(jīng)驗(yàn)值是BCE DiceLoss權(quán)重取 0.5 對(duì) 0.5。BCE 負(fù)責(zé)給所有像素提供穩(wěn)定的梯度信號(hào)DiceLoss 負(fù)責(zé)把前景區(qū)域的梯度變陡。訓(xùn)練初期 BCE 主導(dǎo)后期 DiceLoss 主導(dǎo)兩者互補(bǔ)收斂速度通常比任何單一損失都要快。另外注意這里不要過早引入 Focal Loss 這類復(fù)雜損失它會(huì)加重訓(xùn)練的不穩(wěn)定性先跑通流程再說。4.2 訓(xùn)練循環(huán)、學(xué)習(xí)率與全流程代碼核心訓(xùn)練循環(huán)我把關(guān)鍵部分寫出來包含梯度裁剪、學(xué)習(xí)率調(diào)度和驗(yàn)證指標(biāo)。import torch import torch.nn as nn def train_one_epoch(model, loader, optimizer, criterion, device, clip1.0): model.train() total_loss 0.0 for images, masks in loader: images, masks images.to(device), masks.to(device) logits model(images) # 輸出 shape: [B, 1, H, W] loss criterion(logits, masks.float().unsqueeze(1)) optimizer.zero_grad() loss.backward() nn.utils.clip_grad_norm_(model.parameters(), clip) # 防梯度爆炸 optimizer.step() total_loss loss.item() return total_loss / len(loader) criterion lambda out, target: \ nn.BCEWithLogitsLoss()(out, target) 0.5 * dice_loss(out, target) optimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-5) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max60)dice_loss是自定義函數(shù)對(duì)輸出先過sigmoid然后計(jì)算1 - (2*交集 平滑項(xiàng))/(并集 平滑項(xiàng))。clip_grad_norm_是分割網(wǎng)絡(luò)容易忽略的一步UNet 解碼器層數(shù)多反向傳播時(shí)梯度有累積效應(yīng)尤其是輸入尺寸較大時(shí)不裁剪的話訓(xùn)練中期容易突然發(fā)散。學(xué)習(xí)率從1e-4開始是經(jīng)驗(yàn)值太高會(huì)出現(xiàn) loss 震蕩太低則收斂過慢。CosineAnnealingLR配合 60 個(gè) epoch 的調(diào)度被驗(yàn)證效果穩(wěn)定。我一般在訓(xùn)練時(shí)同時(shí)保留兩份 checkpoint一份按驗(yàn)證 IoU 最高的一份按最新的避免過擬合導(dǎo)致最佳模型丟失。4.3 顯存不足和批大小的心得顯存不足是新人第一個(gè)會(huì)撞到的墻。多數(shù)入門顯卡 6GB 顯存跑 512×512 輸入batch size 只能設(shè)到 2 甚至 1這不是很好但也能跑。優(yōu)先記住先確認(rèn)輸入尺寸是否真的需要 512如果是 256 就能覆蓋目標(biāo)那就用 256。真正對(duì)顯存影響最大的不是輸入的長(zhǎng)寬而是跳躍連接里的特征圖4 個(gè)尺度的特征圖會(huì)在拼接時(shí)同時(shí)駐留顯存。減少初始通道數(shù)可以從 32 降到 16顯存立刻能省三分之一左右。訓(xùn)練時(shí)使用混合精度會(huì)再省一筆常見做法是在 PyTorch 里用自動(dòng)混合精度。另一個(gè)被很多人忽略的點(diǎn)是驗(yàn)證階段的顯存峰值。訓(xùn)練時(shí)梯度是逐層釋放的驗(yàn)證時(shí)不需要梯度但整張?zhí)卣鲌D會(huì)保留到最后所以驗(yàn)證時(shí)的顯存占用反而比訓(xùn)練時(shí)更穩(wěn)定。如果驗(yàn)證時(shí)報(bào) OOM說明是你的輸入尺寸或通道數(shù)實(shí)在太高要回頭降配置。5. 避坑排查UNet 使用時(shí)的常見問題5.1 掩膜尺寸和原圖對(duì)不上現(xiàn)象訓(xùn)練 loss 能下降但輸出圖明顯比輸入小一圈或者有一些奇怪的黑色邊緣。原因數(shù)據(jù)集的長(zhǎng)寬不是固定比例被直接強(qiáng)行 resize 到正方形掩膜和原圖的內(nèi)容發(fā)生了非線性縮放語(yǔ)義錯(cuò)位。另一類是數(shù)據(jù)增強(qiáng)時(shí) train 和 val 用的尺寸不一致評(píng)估時(shí)輸入被打回原始尺寸導(dǎo)致輸出對(duì)不上。解決統(tǒng)一在數(shù)據(jù)加載階段用同一個(gè)(height, width)參數(shù)強(qiáng)轉(zhuǎn)train和val的尺寸保持一致的配置。圖像 resize 用cv2.INTER_LINEAR掩膜用cv2.INTER_NEAREST這兩者配錯(cuò)導(dǎo)致的內(nèi)容錯(cuò)位從訓(xùn)練日志上是看不出來的只有在可視化掩膜覆蓋在原圖上時(shí)才會(huì)暴露。5.2 類別不均衡導(dǎo)致背景“吃掉”前景現(xiàn)象驗(yàn)證集上 IoU 很高比如 0.95但預(yù)測(cè)圖全黑或者小目標(biāo)全丟。原因背景像素占 99%前景占 1%模型學(xué)到了一個(gè)“預(yù)測(cè)全背景”的捷徑。此時(shí)CrossEntropyLoss的數(shù)值很小因?yàn)榻^大多數(shù)像素都被正確分類了但 IoU 指標(biāo)因?yàn)榍熬盀榭斩苯幼兂?0。解決換成BCEWithLogitsLoss加 DiceLoss 的組合給前景類別加權(quán)。另一個(gè)辦法是裁剪訓(xùn)練區(qū)域只看包含前景的局部切片再把切片的坐標(biāo)記錄下來預(yù)處理階段把比例拉均衡。我檢查這個(gè)問題的習(xí)慣是訓(xùn)練到一半時(shí)打印一遍每張預(yù)測(cè)圖的 unique 值如果超過九成都是同一個(gè)值基本就是類別失衡不是模型沒學(xué)好。5.3 學(xué)習(xí)率設(shè)置不當(dāng)訓(xùn)練發(fā)散現(xiàn)象loss 在訓(xùn)練前幾個(gè) batch 正常下降到第 20 個(gè) epoch 附近突然飆到 10 倍以上之后很難恢復(fù)。原因?qū)W習(xí)率太高或者余弦退火降得太快讓權(quán)重更新跨過了局部最優(yōu)而且 UNet 的跳躍連接會(huì)放大這種不穩(wěn)定。解決把lr從1e-3降到1e-4并加上nn.utils.clip_grad_norm_(model.parameters(), 1.0)。更穩(wěn)的做法是每 5 個(gè) epoch 保存一次模型備份訓(xùn)練結(jié)束后選擇驗(yàn)證指標(biāo)最佳的那份而不是選擇最后一份。不要過度依賴早停分割任務(wù)的驗(yàn)證指標(biāo)波動(dòng)比分類任務(wù)大很多早停閾值設(shè)置過小會(huì)把還有上升空間的模型直接砍掉。5.4 數(shù)據(jù)增強(qiáng)把掩膜“旋轉(zhuǎn)錯(cuò)位”之后模型肉眼看不太出來但指標(biāo)不變現(xiàn)象訓(xùn)練 loss 下降正常驗(yàn)證 IoU 停滯不前多個(gè) epoch 都一樣。原因圖像做了旋轉(zhuǎn)而掩膜沒有或者用了不同的隨機(jī)種子導(dǎo)致模型始終學(xué)到“把分割結(jié)果偏移某個(gè)角度”的錯(cuò)位映射。肉眼檢查單張圖有時(shí)候看不出問題這里的關(guān)鍵是統(tǒng)計(jì)驗(yàn)證集所有樣本的 IoU如果多類別任務(wù)某些類別 IoU 一直為 0 但整體指標(biāo)尚可說明增強(qiáng)和驗(yàn)證邏輯之間存在不一致。解決統(tǒng)一用成對(duì)變換作為修復(fù)手段我在__getitem__里把掩膜直接在數(shù)據(jù)加載階段打印出來和原圖疊加可視化人工過一遍確認(rèn)每個(gè)通道都逐像素對(duì)齊后再開始訓(xùn)練。這類問題和大多數(shù)框架報(bào)錯(cuò)無關(guān)屬于“沉默失敗”依賴可視化檢查屬于血淚教訓(xùn)。6. 推理與可視化把模型輸出變成能看的結(jié)果模型訓(xùn)練完只是第一步真正要交付的是“輸入任意一張圖 → 輸出帶標(biāo)注的圖”。推理階段的核心是把模型的 logits 轉(zhuǎn)成類別掩膜再疊加回原圖。這里有一個(gè)容易出問題的細(xì)節(jié)預(yù)測(cè)時(shí)要保持和訓(xùn)練時(shí)相同的預(yù)處理流程尤其是mean/std歸一化和 resize 尺寸否則預(yù)測(cè)效果會(huì)有肉眼可見的下降。我常用的推理代碼是把模型的sigmoid輸出以 0.5 為閾值生成二值掩膜然后用半透明的紅色在原圖上疊加再用opencv的addWeighted混合。保存結(jié)果時(shí)用matplotlib或cv2.imwrite直接存成 PNG。如果同時(shí)需要計(jì)算 IoU要記得把掩膜和標(biāo)簽都轉(zhuǎn)成torch.int64并壓縮到[0, 1]避免數(shù)組廣播帶來的尺寸錯(cuò)位問題。整套流程跑通之后還有兩件事值得做。一件是把模型導(dǎo)出成 TorchScript 或 ONNX這樣可以在移動(dòng)端和邊緣設(shè)備上做推理省去模型架構(gòu)和預(yù)處理代碼的重復(fù)編寫開銷我的經(jīng)驗(yàn)是 ONNX 導(dǎo)出最省心直接用torch.onnx.export配合固定尺寸輸入就能一步到位。另一件事是把數(shù)據(jù)增強(qiáng)、訓(xùn)練和推理腳本封裝成一個(gè)帶命令行參數(shù)的main.py按數(shù)據(jù)路徑和模型權(quán)重兩個(gè)參數(shù)就能調(diào)起整個(gè)流程這不算復(fù)雜工程但對(duì)復(fù)現(xiàn)和后續(xù)修改的幫助非常大。從那以后我每次做分割項(xiàng)目都會(huì)在訓(xùn)練前強(qiáng)制走一遍“數(shù)據(jù)可視化 → 跑通 1 個(gè) batch → 再加增強(qiáng)”的流程這三個(gè)步驟能攔截掉大部分靜默 bug。希望這篇實(shí)戰(zhàn)筆記對(duì)你跑通 UNet 圖像分割有幫助。本文還有配套的精品資源點(diǎn)擊獲取