簡介Unet分割實戰項目面向深度學習圖像分割入門與進階讀者提供DUT-OMRON二值圖像分割數據集及完整訓練/預測代碼。數據集包含4135張訓練圖片及其mask、1033張測試圖片及其mask訓練集與測試集均含images和masks兩個目錄便于對照加載可支撐二分類分割任務。項目實現多尺度訓練train腳本自動計算灰度均值/方差用于歸一化隨機縮放至設定尺寸的0.51.5倍預處理函數集中在transforms.py便于二次修改。網絡訓練50個epochmIoU約0.72采用cos學習率衰減run_results目錄保存損失與IoU曲線、訓練日志和最優權重預測腳本可一鍵推理inference下全部圖片。壓縮包共2000個文件以1979個png圖像/掩碼為主體輔以9個py腳本、5個xml配置、2個txt說明及readme總體約223.63MB。已有441人學習下載適合直接跑通Unet分割流程并遷移至自定義數據也可作為課程設計或論文實驗的基線項目。1. 為什么用 Unet 做 DUT-OMRON 二值分割DUT-OMRON 是顯著性物體檢測中最常被拿來跑深度分割模型的數據集之一單張圖片里“像素屬于顯著前景還是背景”這個問題本質就是二值圖像分割。很多人一上來就套語義分割框架用帶背景類的多分類損失結果發現模型收斂慢、邊緣糊原因不是模型錯而是任務定義沒對準。Unet 的優勢在于它能把空間細節從淺層直接傳遞到深層輸出對像頭發絲、物體邊界這種“細節敏感”場景特別穩。這篇文章不會貼一個已經訓練好的成果讓你膜拜而是沿著“數據 → 網絡 → 訓練 → 推理”的路線把 DUT-OMRON 上用 Unet 做二值分割的完整套路講清楚適合已經會跑通分類任務、想第一次碰分割的讀者也適合想快速對比各種改進 Unet 的工程師。2. Unet 結構拆解與 DUT-OMRON 二值分割的任務定義先明確任務邊界。DUT-OMRON 的標注是像素級二值掩膜白色表示顯著物體黑色表示背景不存在第三個類別。雖然它屬于顯著性物體檢測領域但直接套用現有的分割流程時把它當作二值分割任務來處理是效率最高的做法。2.1 二值分割不是“兩類分類”是逐像素二分類分類模型輸出的是一張圖片屬于貓或狗的概率分割模型輸出的則是和輸入尺寸一致的二維概率圖。對于二值分割輸出通道數為 1每個像素經過 Sigmoid 后就得到一個 0-1 的概率值。訓練時最直接的損失是 BCEWithLogitsLoss它把 Sigmoid 和 BCE 合并在一起計算數值上更穩定。celoss nn.BCEWithLogitsLoss() loss celoss(pred, target)只使用 BCE 的問題在于DUT-OMRON 中背景像素占比經常超過 70%。模型只要學會把所有像素輸出成背景BCE 也能得到一個看起來不錯的數值但分割結果全黑。解決方法是組合 Dice Loss它直接衡量預測掩膜和真實掩膜的重疊度對前景小、背景大的情況天然不敏感。一個簡單的實現如下。def dice_loss(pred_probs, target, smooth1.0): # pred_probs 是經過 sigmoid 的概率target 是二值掩膜 inter (pred_probs * target).sum(dim(2, 3)) union pred_probs.sum(dim(2, 3)) target.sum(dim(2, 3)) dice (2.0 * inter smooth) / (union smooth) return 1.0 - dice.mean()實際訓練時可以取bce 0.5 * dice這里系數 0.5 不是死的先跑 10 個 epoch 看 Dice 變化再調整。如果發現 mask 偏保守就把 dice 權重調大一點如果模型太激進背景上出現大量亮塊就把 dice 權重調小。2.2 Unet 的跳躍連接到底在解決什么Unet 的編碼器會逐步減小分辨率換取更大的感受野最后編碼器輸出的特征圖雖然語義豐富但空間位置已經非常粗糙。解碼器把粗糙的特征圖上采樣回原尺寸此時細節只能靠插值結果就是邊界發糊。跳躍連接的作用是把編碼器每一層下采樣前的高分辨率特征復制一份在解碼器上采樣的同時拼接到同一層級的特征圖上。這樣解碼器既能看到語義又能看到原始邊緣信息。在 DUT-OMRON 上顯著物體往往邊緣復雜例如桌腿、樹枝、發絲與背景重疊。沒有跳躍連接時模型容易把與前景紋理相近的背景區域一起標成前景。有了跳躍連接淺層特征攜帶的梯度能直接反向傳播到編碼器早期讓邊緣位置的權重更快糾正。這也是 Unet 在只有幾千張訓練數據時仍然能訓練出不錯效果的重要原因因為它比普通 FCN 更少依賴訓練樣本量。2.3 DUT-OMRON 數據集格式與讀取策略原始數據集通常直接解壓后就能看到 images 和 masks 兩個目錄但文件名后綴可能不一樣有些掩膜是 bmp有些是 png需要先統一格式。我一般會按下面的目錄結構重新整理目錄或文件存放內容格式要求images/train訓練原圖RGB JPG/PNGmasks/train訓練掩膜單通道 PNG0/255images/val驗證原圖與訓練一致masks/val驗證掩膜與訓練一致train.txt訓練文件名每行一個不帶擴展名val.txt驗證文件名同上讀取掩膜時不要直接cv2.imread(path, cv2.IMREAD_COLOR)因為三通道 RGB 得到的掩膜每個通道一樣但邊緣會有抗鋸齒導致的灰色過渡。正確做法是使用灰度模式讀取然后施加閾值mask cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) mask (mask 127).astype(np.uint8)原圖尺寸如果大小不一需要統一縮放一般最小邊縮到 256 或 320。掩膜縮放必須用最近鄰插值避免線性插值在邊緣產生中間灰值。這一步很多初學者容易忽略導致訓練時損失來回震蕩。3. 用 PyTorch 搭建 Unet 并訓練 DUT-OMRON 數據集理論清楚后接下來講如何用 PyTorch 把訓練流程跑通。這一部分會給出數據集封裝、網絡定義和訓練循環三塊代碼你可以直接復制到自己的項目里改路徑跑腳本。3.1 自定義 Dataset把 DUT-OMRON 圖片和掩膜對齊先準備train.txt和val.txt文件每行只寫圖片名不帶擴展名。然后實現 Dataset 類import cv2 import torch import numpy as np from torch.utils.data import Dataset class DUTSegDataset(Dataset): def __init__(self, img_dir, mask_dir, file_list, size256): with open(file_list, r) as f: self.names [line.strip() for line in f.readlines() if line.strip()] self.img_dir img_dir self.mask_dir mask_dir self.size size def __len__(self): return len(self.names) def _load_images(self, name): img cv2.imread(self.img_dir / name .jpg) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) mask cv2.imread(self.mask_dir / name .png, cv2.IMREAD_GRAYSCALE) if mask is None: raise FileNotFoundError(fMask not found for {name}) mask (mask 127).astype(np.uint8) return img, mask def __getitem__(self, idx): name self.names[idx] img, mask self._load_images(name) img cv2.resize(img, (self.size, self.size), interpolationcv2.INTER_LINEAR) mask cv2.resize(mask, (self.size, self.size), interpolationcv2.INTER_NEAREST) img img.astype(np.float32) / 255.0 img torch.from_numpy(img).permute(2, 0, 1) mask torch.from_numpy(mask.astype(np.float32)).unsqueeze(0) return img, mask這里需要注意三點一是圖片讀入后經過了BGR2RGB轉換和訓練的通道順序保持一致二是掩膜在讀入后立即做了閾值化切斷灰邊影響三是縮放時原圖用線性插值掩膜用最近鄰。如果數據集圖片本身就固定是 256x256可以跳過 resize但保留這一行對換數據集參測時更省心。3.2 輕量 Unet 的 PyTorch 實現這里給一個可以直接跑通實驗的輕量版本。它保留了 Unet 最核心的雙卷積塊、下采樣、轉置卷積和跳躍連接結構適合在單張消費級顯卡上調試 DUT-OMRON。import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), ) def forward(self, x): return self.conv(x) class LightUNet(nn.Module): def __init__(self, in_ch3, out_ch1): super().__init__() self.enc1 DoubleConv(in_ch, 64) self.enc2 DoubleConv(64, 128) self.enc3 DoubleConv(128, 256) self.pool nn.MaxPool2d(2) self.center DoubleConv(256, 512) self.up2 nn.ConvTranspose2d(512, 256, 2, stride2) self.dec2 DoubleConv(512, 256) self.up1 nn.ConvTranspose2d(256, 128, 2, stride2) self.dec1 DoubleConv(256, 128) self.up0 nn.ConvTranspose2d(128, 64, 2, stride2) self.dec0 DoubleConv(128, 64) self.out nn.Conv2d(64, out_ch, 1) def forward(self, x): e1 self.enc1(x) e2 self.enc2(self.pool(e1)) e3 self.enc3(self.pool(e2)) c self.center(self.pool(e3)) d2 self.dec2(torch.cat([self.up2(c), e3], dim1)) d1 self.dec1(torch.cat([self.up1(d2), e2], dim1)) d0 self.dec0(torch.cat([self.up0(d1), e1], dim1)) return self.out(d0)注意這個結構只做了三次下采樣輸入 256x256 時最底層特征圖是 32x32參數量比標準 Unet 小很多訓練速度快適合驗證流程是否通順。要把網絡換成標準 Unet只需要把每個模塊的通道數改成 64、128、256、512再加一層下采樣即可。3.3 訓練循環與驗證流程訓練流程主體是一個 for 循環里面包含前向傳播、損失計算、反向傳播和驗證。為節省篇幅這里只寫每輪的核心步驟。import torch.optim as optim from torch.utils.data import DataLoader from torch.cuda.amp import GradScaler, autocast model LightUNet(3, 1).cuda() dataset DUTSegDataset(images/train, masks/train, train.txt, size256) loader DataLoader(dataset, batch_size16, shuffleTrue, num_workers4, drop_lastTrue) optimizer optim.AdamW(model.parameters(), lr1e-3, weight_decay1e-4) scaler GradScaler() celoss nn.BCEWithLogitsLoss() for epoch in range(30): model.train() for imgs, masks in loader: imgs, masks imgs.cuda(), masks.cuda() optimizer.zero_grad() with autocast(): logits model(imgs) loss celoss(logits, masks) 0.5 * dice_loss(torch.sigmoid(logits), masks) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()autocast和GradScaler是 PyTorch 混合精度訓練的標準組合只要顯卡支持半精度就能帶來明顯速度提升。Batch Size 按顯存調整我一般固定輸入 256x256 時開 16 或 8。如果顯存只有 6Gbatch size 降到 4同時把模型第一層通道數從 64 改成 32。超參數推薦范圍說明輸入尺寸256x256 / 320x320顯存小選 224batch size8-16依賴 GPU 顯存優化器AdamW/AdamAdam 收斂快AdamW 泛化稍好學習率1e-4 到 3e-3用余弦退火時從 3e-3 起步epochs30-60DUT-OMRON 樣本量50 輪足夠提示訓練時如果 GPU 顯存不足優先減小 batch size而不要降低輸入尺寸因為 256 和 224 的精度差距在 DUT-OMRON 上能直接反映到 F-measure 上。這段代碼可以作為基線。把train.txt和val.txt換成自己的劃分文件就能直接跑起來。如果追求更高精度通常更快的路徑是在預訓練的 ResNet34 encoder 上做 Unet但那是下一步的改進方向。4. 訓練 DUT-OMRON 時的關鍵參數與常見坑前面的代碼可以跑但真正能把 Unet 在 DUT-OMRON 上訓練到實用水平還需要處理好四個參數相關的問題。這一章把訓練過程中最常遇到的坑挑出來對應的調整路徑也一起給出。4.1 圖像尺寸、batch size 與顯存的權衡如果直接使用 512x512 的原圖訓練顯存消耗會以平方級別增長可能一個 batch 都放不下。DUT-OMRON 的很多圖片背景復雜但顯著物體通常不明顯所以一開始不需要上大分辨率。常見做法是先用 256x256 跑通全流程驗證模型能收斂后再嘗試 384 或 512 微調。切換分辨率時要注意兩點一是輸入尺寸必須是 16 的整數倍否則多次下采樣后特征圖尺寸對不上二是掩膜 resize 必須固定使用最近鄰插值這是老生常談但仍然經常出錯的點。4.2 學習率與優化器選擇對于二值分割任務AdamW 比 SGD 更省心因為它自帶了動量和自適應學習率適合初學者第一次訓練時不必手動調太多參數。推薦初始學習率設為1e-3配合余弦退火或 ReduceLROnPlateau 調度器在驗證集不漲時降學習率。下面這段配置是我常用的optimizer optim.AdamW(model.parameters(), lr1e-3, weight_decay1e-4) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max30)如果使用 CosineAnnealingLRT_max 通常設為總 epoch 數。學習率過高的癥狀是 loss 在初期下降后立刻反彈說明前幾步跨過了最優區域學習率過低則損失曲線下降像蝸牛。訓練初期應觀察前 3 個 epoch 的訓練 loss 是否明顯下降如果不降就調大學習率試一次。4.3 類別不均衡與損失函數權重DUT-OMRON 中的顯著物體沒有固定大小有些圖前景占比不到 10%。這樣的樣本如果只用 BCE反向傳播時背景像素貢獻的梯度會淹沒前景。上一章的組合損失可以緩解但還需要做一些額外處理。一種很有效的技巧是對前景像素做指數懲罰例如使用 Tversky Loss 或 Focal Loss。以 Focal Loss 為例它對難分類的像素分配更高權重能明顯提升小物體召回率。def focal_loss(logits, targets, gamma2.0, alpha0.25): probs torch.sigmoid(logits) ce -targets * torch.log(probs 1e-6) - (1 - targets) * torch.log(1 - probs 1e-6) pt targets * probs (1 - targets) * (1 - probs) return (alpha * (1 - pt) ** gamma * ce).mean()不過 Focal Loss 的超參數gamma、alpha需要在驗證集上多試幾次。經驗值是gamma2、alpha0.25對前景小的場景有效如果前景占比本身不小直接用 BCEDice 就夠了。4.4 過擬合與驗證集劃分DUT-OMRON 雖然有五千多張圖但很多圖片的掩膜存在大量空白背景模型容易在訓練集上死記背景紋理。驗證集至少要留 300 張且劃分時要保證顯著物體大小分布相似不要讓驗證集里全是全黑掩膜。我習慣將數據集按文件名哈希值劃分出約 15% 作為驗證集。現象可能原因處理方式訓練 loss 下降但驗證 loss 上升過擬合增加數據增強、減小模型容量、增加 weight decay驗證集出現全黑預測前景像素占比太低后處理閾值太高降低預測閾值到 0.3再觀察邊緣區域出現灰帶掩膜 resize 使用了線性插值改成最近鄰插值重新生成數據模型前期 loss 不降學習率過大或輸入未歸一化調小學習率檢查像素值范圍數據增強也是反過擬合的重要一環。最安全的增強包括隨機水平翻轉、隨機小范圍旋轉和隨機亮度擾動。對于分割任務旋轉角度要控制在 15° 以內翻轉時原圖和掩膜必須用同樣的隨機種子同步變化。5. 用訓練好的 Unet 做推理與后處理訓練完模型最關鍵的一步是把模型輸出轉換成可展示或可交付的二值圖。這一章給出完整的推理流程包括模型加載、閾值選擇、后處理和常見形狀錯誤排查。5.1 加載模型權重并切換到 eval 模式訓練結束后保存模型有兩種常見方式存整個模型和只存 state_dict。推薦后者因為它體積小且不受 PyTorch 版本內部結構影響。推理前需要先實例化同樣的網絡結構然后加載權重model LightUNet(3, 1) state torch.load(best_model.pth, map_locationcpu) model.load_state_dict(state) model.eval().cuda()注意model.eval()會關閉 Dropout 和 BatchNorm 的統計量更新邏輯。如果漏掉這一步BatchNorm 層會因為使用訓練時的 batch 統計而讓預測結果出現隨機抖動。5.2 單張圖像預測與閾值化推理時輸入圖像需要和訓練時做完全一樣的預處理讀圖、BGR2RGB、縮放、歸一化到 0-1。預測輸出的 logits 尺寸是[1, 1, H, W]經過 Sigmoid 后得到概率圖。二值分割通常直接用 0.5 作為閾值但建議在驗證集上嘗試 0.3 到 0.7 之間步長 0.05 的閾值選擇 F-measure 最高的那一個。import cv2 import numpy as np import torch def prepare_image(img_path, size256): img cv2.imread(img_path) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img cv2.resize(img, (size, size), interpolationcv2.INTER_LINEAR) img img.astype(np.float32) / 255.0 return torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0) with torch.no_grad(): img_input prepare_image(img_path, size256).cuda() logits model(img_input) probs torch.sigmoid(logits).cpu().numpy().squeeze() mask (probs 0.5).astype(np.uint8)這里的prepare_image重點是把像素縮放到 0-1 并且轉換為[C, H, W]的展平輸入。如果原圖不是正方形縮放后會和訓練時的分辨率不一致所以訓練和推理最好統一固定尺寸比如都縮放到 256x256再做后處理時再把掩膜 resize 回原圖像大小。5.3 后處理去除小連通域和填充孔洞模型輸出的二值圖常常帶一些孤立的小亮點原因是背景中某些紋理被誤判為前景。常見做法是使用連通域分析把面積小于閾值的連通域去掉。連通域的面積閾值一般按圖片尺寸比例設定例如0.5% * H * W。import cv2 import numpy as np def remove_small_areas(mask, min_area_ratio0.005): num_labels, labels, stats, _ cv2.connectedComponentsWithStats(mask.astype(np.uint8), 8) out np.zeros_like(mask) for label in range(1, num_labels): if stats[label, cv2.CC_STAT_AREA] min_area_ratio * mask.size: out[labels label] 1 return out另外如果掩膜內部出現小孔可以使用cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)做一次閉運算。但要注意閉運算也會讓原本細長的邊界變厚所以對于邊緣精度要求高的場景閉運算的 kernel 要盡量小比如 3x3。5.4 常見推理錯誤排查推理時的報錯和訓練時不一樣下面列出幾個最容易碰到的錯誤現象報錯關鍵詞解決思路輸出尺寸與原圖不一致size mismatch檢查 resize 是否發生在模型之前GPU 顯存溢出CUDA out of memory降低 batch size 到 1輸入尺寸調小預測結果全一樣all same檢查是否忘了model.eval()掩膜邊緣有灰色過渡value not 0/1檢查后處理時是否做了閾值化遇到顯存溢出時也可以用半精度推理把模型參數轉成half()輸入也轉成half()在大多數顯卡上能直接省接近一半顯存。6. Unet 改進方向與驗證技巧當基線跑通后再往前走通常有兩個方向一是改網絡結構讓邊界更精細二是改訓練目標讓顯著物體更完整。這里分享兩個實戰中有效的技巧可以直接套在 DUT-OMRON 上驗證。6.1 Deep Supervision給解碼器每個分支加輔助損失標準 Unet 只在最后一層計算損失梯度需要從很深的網絡穿越回編碼器。Deep Supervision 的思想是在解碼器的每個上采樣層級都接一個 1x1 卷積和損失計算然后把多個損失加權求和。這樣編碼器每一層都能直接拿到監督信號收斂速度快尤其在數據量只有幾千張時效果更明顯。實現很簡單在訓練時讓LightUNet的forward函數返回一個列表[out1, out2, out3]分別對應三個解碼層輸出然后對每個輸出算一次 BCEDice再按系數加起來。需要關注的是各輔助損失在總損失中的權重一般從 0.3 開始主分支權重保持 1.0。6.2 用 F-measure 驗證而不是只用 Accuracy二值分割任務里Accuracy 會被大量背景像素拉高即使把整個圖預測成黑色也能得到 80% 以上的精度。DUT-OMRON 這種顯著性分割場景最常用的評價指標是 F-measure它是精確率和召回率的調和平均能同時懲罰漏檢和誤檢?？梢栽隍炞C循環里加上這段計算precision (pred target).sum() / (pred.sum() 1e-6) recall (pred target).sum() / (target.sum() 1e-6) f_measure 2 * precision * recall / (precision recall 1e-6)需要根據驗證集上的 F-measure 來選擇后處理閾值。常見的做法是在驗證集上對每個候選閾值計算 F-measure取平均值最高的閾值作為最終閾值再對測試集做推理。這套流程雖然簡單但比固定 0.5 更貼合數據分布。6.3 從 256 到原分辨率輸出 mask 的恢復如果訓練時把原圖縮放到 256x256推理得到 256x256 的掩膜后需要把它放回原圖尺寸參與指標計算。這里有一個很容易被忽略的點cv2.resize對 mask 要用最鄰近插值但如果你用 PIL 的Image.resize默認使用 BILINEAR會產生非 0/1 的中間值。務必要顯式指定Image.NEAREST。所有掩膜后處理操作建議統一在 CPU 上用 OpenCV 完成不要用 GPU 張量直接 resize因為 PyTorch 的 interpolate 默認是雙線性輸出會變成浮點掩膜。保存結果時再強制mask[mask 0.5] 1; mask[mask 0.5] 0。這樣交付的掩膜才是純 0/1不會因為格式問題被下游指標腳本誤判。本文還有配套的精品資源點擊獲取 用Unet實現DUT-OMRON二值分割：從數據到訓練的PyTorch實戰
簡介Unet分割實戰項目面向深度學習圖像分割入門與進階讀者提供DUT-OMRON二值圖像分割數據集及完整訓練/預測代碼。數據集包含4135張訓練圖片及其mask、1033張測試圖片及其mask訓練集與測試集均含images和masks兩個目錄便于對照加載可支撐二分類分割任務。項目實現多尺度訓練train腳本自動計算灰度均值/方差用于歸一化隨機縮放至設定尺寸的0.51.5倍預處理函數集中在transforms.py便于二次修改。網絡訓練50個epochmIoU約0.72采用cos學習率衰減run_results目錄保存損失與IoU曲線、訓練日志和最優權重預測腳本可一鍵推理inference下全部圖片。壓縮包共2000個文件以1979個png圖像/掩碼為主體輔以9個py腳本、5個xml配置、2個txt說明及readme總體約223.63MB。已有441人學習下載適合直接跑通Unet分割流程并遷移至自定義數據也可作為課程設計或論文實驗的基線項目。1. 為什么用 Unet 做 DUT-OMRON 二值分割DUT-OMRON 是顯著性物體檢測中最常被拿來跑深度分割模型的數據集之一單張圖片里“像素屬于顯著前景還是背景”這個問題本質就是二值圖像分割。很多人一上來就套語義分割框架用帶背景類的多分類損失結果發現模型收斂慢、邊緣糊原因不是模型錯而是任務定義沒對準。Unet 的優勢在于它能把空間細節從淺層直接傳遞到深層輸出對像頭發絲、物體邊界這種“細節敏感”場景特別穩。這篇文章不會貼一個已經訓練好的成果讓你膜拜而是沿著“數據 → 網絡 → 訓練 → 推理”的路線把 DUT-OMRON 上用 Unet 做二值分割的完整套路講清楚適合已經會跑通分類任務、想第一次碰分割的讀者也適合想快速對比各種改進 Unet 的工程師。2. Unet 結構拆解與 DUT-OMRON 二值分割的任務定義先明確任務邊界。DUT-OMRON 的標注是像素級二值掩膜白色表示顯著物體黑色表示背景不存在第三個類別。雖然它屬于顯著性物體檢測領域但直接套用現有的分割流程時把它當作二值分割任務來處理是效率最高的做法。2.1 二值分割不是“兩類分類”是逐像素二分類分類模型輸出的是一張圖片屬于貓或狗的概率分割模型輸出的則是和輸入尺寸一致的二維概率圖。對于二值分割輸出通道數為 1每個像素經過 Sigmoid 后就得到一個 0-1 的概率值。訓練時最直接的損失是 BCEWithLogitsLoss它把 Sigmoid 和 BCE 合并在一起計算數值上更穩定。celoss nn.BCEWithLogitsLoss() loss celoss(pred, target)只使用 BCE 的問題在于DUT-OMRON 中背景像素占比經常超過 70%。模型只要學會把所有像素輸出成背景BCE 也能得到一個看起來不錯的數值但分割結果全黑。解決方法是組合 Dice Loss它直接衡量預測掩膜和真實掩膜的重疊度對前景小、背景大的情況天然不敏感。一個簡單的實現如下。def dice_loss(pred_probs, target, smooth1.0): # pred_probs 是經過 sigmoid 的概率target 是二值掩膜 inter (pred_probs * target).sum(dim(2, 3)) union pred_probs.sum(dim(2, 3)) target.sum(dim(2, 3)) dice (2.0 * inter smooth) / (union smooth) return 1.0 - dice.mean()實際訓練時可以取bce 0.5 * dice這里系數 0.5 不是死的先跑 10 個 epoch 看 Dice 變化再調整。如果發現 mask 偏保守就把 dice 權重調大一點如果模型太激進背景上出現大量亮塊就把 dice 權重調小。2.2 Unet 的跳躍連接到底在解決什么Unet 的編碼器會逐步減小分辨率換取更大的感受野最后編碼器輸出的特征圖雖然語義豐富但空間位置已經非常粗糙。解碼器把粗糙的特征圖上采樣回原尺寸此時細節只能靠插值結果就是邊界發糊。跳躍連接的作用是把編碼器每一層下采樣前的高分辨率特征復制一份在解碼器上采樣的同時拼接到同一層級的特征圖上。這樣解碼器既能看到語義又能看到原始邊緣信息。在 DUT-OMRON 上顯著物體往往邊緣復雜例如桌腿、樹枝、發絲與背景重疊。沒有跳躍連接時模型容易把與前景紋理相近的背景區域一起標成前景。有了跳躍連接淺層特征攜帶的梯度能直接反向傳播到編碼器早期讓邊緣位置的權重更快糾正。這也是 Unet 在只有幾千張訓練數據時仍然能訓練出不錯效果的重要原因因為它比普通 FCN 更少依賴訓練樣本量。2.3 DUT-OMRON 數據集格式與讀取策略原始數據集通常直接解壓后就能看到 images 和 masks 兩個目錄但文件名后綴可能不一樣有些掩膜是 bmp有些是 png需要先統一格式。我一般會按下面的目錄結構重新整理目錄或文件存放內容格式要求images/train訓練原圖RGB JPG/PNGmasks/train訓練掩膜單通道 PNG0/255images/val驗證原圖與訓練一致masks/val驗證掩膜與訓練一致train.txt訓練文件名每行一個不帶擴展名val.txt驗證文件名同上讀取掩膜時不要直接cv2.imread(path, cv2.IMREAD_COLOR)因為三通道 RGB 得到的掩膜每個通道一樣但邊緣會有抗鋸齒導致的灰色過渡。正確做法是使用灰度模式讀取然后施加閾值mask cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) mask (mask 127).astype(np.uint8)原圖尺寸如果大小不一需要統一縮放一般最小邊縮到 256 或 320。掩膜縮放必須用最近鄰插值避免線性插值在邊緣產生中間灰值。這一步很多初學者容易忽略導致訓練時損失來回震蕩。3. 用 PyTorch 搭建 Unet 并訓練 DUT-OMRON 數據集理論清楚后接下來講如何用 PyTorch 把訓練流程跑通。這一部分會給出數據集封裝、網絡定義和訓練循環三塊代碼你可以直接復制到自己的項目里改路徑跑腳本。3.1 自定義 Dataset把 DUT-OMRON 圖片和掩膜對齊先準備train.txt和val.txt文件每行只寫圖片名不帶擴展名。然后實現 Dataset 類import cv2 import torch import numpy as np from torch.utils.data import Dataset class DUTSegDataset(Dataset): def __init__(self, img_dir, mask_dir, file_list, size256): with open(file_list, r) as f: self.names [line.strip() for line in f.readlines() if line.strip()] self.img_dir img_dir self.mask_dir mask_dir self.size size def __len__(self): return len(self.names) def _load_images(self, name): img cv2.imread(self.img_dir / name .jpg) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) mask cv2.imread(self.mask_dir / name .png, cv2.IMREAD_GRAYSCALE) if mask is None: raise FileNotFoundError(fMask not found for {name}) mask (mask 127).astype(np.uint8) return img, mask def __getitem__(self, idx): name self.names[idx] img, mask self._load_images(name) img cv2.resize(img, (self.size, self.size), interpolationcv2.INTER_LINEAR) mask cv2.resize(mask, (self.size, self.size), interpolationcv2.INTER_NEAREST) img img.astype(np.float32) / 255.0 img torch.from_numpy(img).permute(2, 0, 1) mask torch.from_numpy(mask.astype(np.float32)).unsqueeze(0) return img, mask這里需要注意三點一是圖片讀入后經過了BGR2RGB轉換和訓練的通道順序保持一致二是掩膜在讀入后立即做了閾值化切斷灰邊影響三是縮放時原圖用線性插值掩膜用最近鄰。如果數據集圖片本身就固定是 256x256可以跳過 resize但保留這一行對換數據集參測時更省心。3.2 輕量 Unet 的 PyTorch 實現這里給一個可以直接跑通實驗的輕量版本。它保留了 Unet 最核心的雙卷積塊、下采樣、轉置卷積和跳躍連接結構適合在單張消費級顯卡上調試 DUT-OMRON。import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), ) def forward(self, x): return self.conv(x) class LightUNet(nn.Module): def __init__(self, in_ch3, out_ch1): super().__init__() self.enc1 DoubleConv(in_ch, 64) self.enc2 DoubleConv(64, 128) self.enc3 DoubleConv(128, 256) self.pool nn.MaxPool2d(2) self.center DoubleConv(256, 512) self.up2 nn.ConvTranspose2d(512, 256, 2, stride2) self.dec2 DoubleConv(512, 256) self.up1 nn.ConvTranspose2d(256, 128, 2, stride2) self.dec1 DoubleConv(256, 128) self.up0 nn.ConvTranspose2d(128, 64, 2, stride2) self.dec0 DoubleConv(128, 64) self.out nn.Conv2d(64, out_ch, 1) def forward(self, x): e1 self.enc1(x) e2 self.enc2(self.pool(e1)) e3 self.enc3(self.pool(e2)) c self.center(self.pool(e3)) d2 self.dec2(torch.cat([self.up2(c), e3], dim1)) d1 self.dec1(torch.cat([self.up1(d2), e2], dim1)) d0 self.dec0(torch.cat([self.up0(d1), e1], dim1)) return self.out(d0)注意這個結構只做了三次下采樣輸入 256x256 時最底層特征圖是 32x32參數量比標準 Unet 小很多訓練速度快適合驗證流程是否通順。要把網絡換成標準 Unet只需要把每個模塊的通道數改成 64、128、256、512再加一層下采樣即可。3.3 訓練循環與驗證流程訓練流程主體是一個 for 循環里面包含前向傳播、損失計算、反向傳播和驗證。為節省篇幅這里只寫每輪的核心步驟。import torch.optim as optim from torch.utils.data import DataLoader from torch.cuda.amp import GradScaler, autocast model LightUNet(3, 1).cuda() dataset DUTSegDataset(images/train, masks/train, train.txt, size256) loader DataLoader(dataset, batch_size16, shuffleTrue, num_workers4, drop_lastTrue) optimizer optim.AdamW(model.parameters(), lr1e-3, weight_decay1e-4) scaler GradScaler() celoss nn.BCEWithLogitsLoss() for epoch in range(30): model.train() for imgs, masks in loader: imgs, masks imgs.cuda(), masks.cuda() optimizer.zero_grad() with autocast(): logits model(imgs) loss celoss(logits, masks) 0.5 * dice_loss(torch.sigmoid(logits), masks) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()autocast和GradScaler是 PyTorch 混合精度訓練的標準組合只要顯卡支持半精度就能帶來明顯速度提升。Batch Size 按顯存調整我一般固定輸入 256x256 時開 16 或 8。如果顯存只有 6Gbatch size 降到 4同時把模型第一層通道數從 64 改成 32。超參數推薦范圍說明輸入尺寸256x256 / 320x320顯存小選 224batch size8-16依賴 GPU 顯存優化器AdamW/AdamAdam 收斂快AdamW 泛化稍好學習率1e-4 到 3e-3用余弦退火時從 3e-3 起步epochs30-60DUT-OMRON 樣本量50 輪足夠提示訓練時如果 GPU 顯存不足優先減小 batch size而不要降低輸入尺寸因為 256 和 224 的精度差距在 DUT-OMRON 上能直接反映到 F-measure 上。這段代碼可以作為基線。把train.txt和val.txt換成自己的劃分文件就能直接跑起來。如果追求更高精度通常更快的路徑是在預訓練的 ResNet34 encoder 上做 Unet但那是下一步的改進方向。4. 訓練 DUT-OMRON 時的關鍵參數與常見坑前面的代碼可以跑但真正能把 Unet 在 DUT-OMRON 上訓練到實用水平還需要處理好四個參數相關的問題。這一章把訓練過程中最常遇到的坑挑出來對應的調整路徑也一起給出。4.1 圖像尺寸、batch size 與顯存的權衡如果直接使用 512x512 的原圖訓練顯存消耗會以平方級別增長可能一個 batch 都放不下。DUT-OMRON 的很多圖片背景復雜但顯著物體通常不明顯所以一開始不需要上大分辨率。常見做法是先用 256x256 跑通全流程驗證模型能收斂后再嘗試 384 或 512 微調。切換分辨率時要注意兩點一是輸入尺寸必須是 16 的整數倍否則多次下采樣后特征圖尺寸對不上二是掩膜 resize 必須固定使用最近鄰插值這是老生常談但仍然經常出錯的點。4.2 學習率與優化器選擇對于二值分割任務AdamW 比 SGD 更省心因為它自帶了動量和自適應學習率適合初學者第一次訓練時不必手動調太多參數。推薦初始學習率設為1e-3配合余弦退火或 ReduceLROnPlateau 調度器在驗證集不漲時降學習率。下面這段配置是我常用的optimizer optim.AdamW(model.parameters(), lr1e-3, weight_decay1e-4) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max30)如果使用 CosineAnnealingLRT_max 通常設為總 epoch 數。學習率過高的癥狀是 loss 在初期下降后立刻反彈說明前幾步跨過了最優區域學習率過低則損失曲線下降像蝸牛。訓練初期應觀察前 3 個 epoch 的訓練 loss 是否明顯下降如果不降就調大學習率試一次。4.3 類別不均衡與損失函數權重DUT-OMRON 中的顯著物體沒有固定大小有些圖前景占比不到 10%。這樣的樣本如果只用 BCE反向傳播時背景像素貢獻的梯度會淹沒前景。上一章的組合損失可以緩解但還需要做一些額外處理。一種很有效的技巧是對前景像素做指數懲罰例如使用 Tversky Loss 或 Focal Loss。以 Focal Loss 為例它對難分類的像素分配更高權重能明顯提升小物體召回率。def focal_loss(logits, targets, gamma2.0, alpha0.25): probs torch.sigmoid(logits) ce -targets * torch.log(probs 1e-6) - (1 - targets) * torch.log(1 - probs 1e-6) pt targets * probs (1 - targets) * (1 - probs) return (alpha * (1 - pt) ** gamma * ce).mean()不過 Focal Loss 的超參數gamma、alpha需要在驗證集上多試幾次。經驗值是gamma2、alpha0.25對前景小的場景有效如果前景占比本身不小直接用 BCEDice 就夠了。4.4 過擬合與驗證集劃分DUT-OMRON 雖然有五千多張圖但很多圖片的掩膜存在大量空白背景模型容易在訓練集上死記背景紋理。驗證集至少要留 300 張且劃分時要保證顯著物體大小分布相似不要讓驗證集里全是全黑掩膜。我習慣將數據集按文件名哈希值劃分出約 15% 作為驗證集。現象可能原因處理方式訓練 loss 下降但驗證 loss 上升過擬合增加數據增強、減小模型容量、增加 weight decay驗證集出現全黑預測前景像素占比太低后處理閾值太高降低預測閾值到 0.3再觀察邊緣區域出現灰帶掩膜 resize 使用了線性插值改成最近鄰插值重新生成數據模型前期 loss 不降學習率過大或輸入未歸一化調小學習率檢查像素值范圍數據增強也是反過擬合的重要一環。最安全的增強包括隨機水平翻轉、隨機小范圍旋轉和隨機亮度擾動。對于分割任務旋轉角度要控制在 15° 以內翻轉時原圖和掩膜必須用同樣的隨機種子同步變化。5. 用訓練好的 Unet 做推理與后處理訓練完模型最關鍵的一步是把模型輸出轉換成可展示或可交付的二值圖。這一章給出完整的推理流程包括模型加載、閾值選擇、后處理和常見形狀錯誤排查。5.1 加載模型權重并切換到 eval 模式訓練結束后保存模型有兩種常見方式存整個模型和只存 state_dict。推薦后者因為它體積小且不受 PyTorch 版本內部結構影響。推理前需要先實例化同樣的網絡結構然后加載權重model LightUNet(3, 1) state torch.load(best_model.pth, map_locationcpu) model.load_state_dict(state) model.eval().cuda()注意model.eval()會關閉 Dropout 和 BatchNorm 的統計量更新邏輯。如果漏掉這一步BatchNorm 層會因為使用訓練時的 batch 統計而讓預測結果出現隨機抖動。5.2 單張圖像預測與閾值化推理時輸入圖像需要和訓練時做完全一樣的預處理讀圖、BGR2RGB、縮放、歸一化到 0-1。預測輸出的 logits 尺寸是[1, 1, H, W]經過 Sigmoid 后得到概率圖。二值分割通常直接用 0.5 作為閾值但建議在驗證集上嘗試 0.3 到 0.7 之間步長 0.05 的閾值選擇 F-measure 最高的那一個。import cv2 import numpy as np import torch def prepare_image(img_path, size256): img cv2.imread(img_path) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img cv2.resize(img, (size, size), interpolationcv2.INTER_LINEAR) img img.astype(np.float32) / 255.0 return torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0) with torch.no_grad(): img_input prepare_image(img_path, size256).cuda() logits model(img_input) probs torch.sigmoid(logits).cpu().numpy().squeeze() mask (probs 0.5).astype(np.uint8)這里的prepare_image重點是把像素縮放到 0-1 并且轉換為[C, H, W]的展平輸入。如果原圖不是正方形縮放后會和訓練時的分辨率不一致所以訓練和推理最好統一固定尺寸比如都縮放到 256x256再做后處理時再把掩膜 resize 回原圖像大小。5.3 后處理去除小連通域和填充孔洞模型輸出的二值圖常常帶一些孤立的小亮點原因是背景中某些紋理被誤判為前景。常見做法是使用連通域分析把面積小于閾值的連通域去掉。連通域的面積閾值一般按圖片尺寸比例設定例如0.5% * H * W。import cv2 import numpy as np def remove_small_areas(mask, min_area_ratio0.005): num_labels, labels, stats, _ cv2.connectedComponentsWithStats(mask.astype(np.uint8), 8) out np.zeros_like(mask) for label in range(1, num_labels): if stats[label, cv2.CC_STAT_AREA] min_area_ratio * mask.size: out[labels label] 1 return out另外如果掩膜內部出現小孔可以使用cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)做一次閉運算。但要注意閉運算也會讓原本細長的邊界變厚所以對于邊緣精度要求高的場景閉運算的 kernel 要盡量小比如 3x3。5.4 常見推理錯誤排查推理時的報錯和訓練時不一樣下面列出幾個最容易碰到的錯誤現象報錯關鍵詞解決思路輸出尺寸與原圖不一致size mismatch檢查 resize 是否發生在模型之前GPU 顯存溢出CUDA out of memory降低 batch size 到 1輸入尺寸調小預測結果全一樣all same檢查是否忘了model.eval()掩膜邊緣有灰色過渡value not 0/1檢查后處理時是否做了閾值化遇到顯存溢出時也可以用半精度推理把模型參數轉成half()輸入也轉成half()在大多數顯卡上能直接省接近一半顯存。6. Unet 改進方向與驗證技巧當基線跑通后再往前走通常有兩個方向一是改網絡結構讓邊界更精細二是改訓練目標讓顯著物體更完整。這里分享兩個實戰中有效的技巧可以直接套在 DUT-OMRON 上驗證。6.1 Deep Supervision給解碼器每個分支加輔助損失標準 Unet 只在最后一層計算損失梯度需要從很深的網絡穿越回編碼器。Deep Supervision 的思想是在解碼器的每個上采樣層級都接一個 1x1 卷積和損失計算然后把多個損失加權求和。這樣編碼器每一層都能直接拿到監督信號收斂速度快尤其在數據量只有幾千張時效果更明顯。實現很簡單在訓練時讓LightUNet的forward函數返回一個列表[out1, out2, out3]分別對應三個解碼層輸出然后對每個輸出算一次 BCEDice再按系數加起來。需要關注的是各輔助損失在總損失中的權重一般從 0.3 開始主分支權重保持 1.0。6.2 用 F-measure 驗證而不是只用 Accuracy二值分割任務里Accuracy 會被大量背景像素拉高即使把整個圖預測成黑色也能得到 80% 以上的精度。DUT-OMRON 這種顯著性分割場景最常用的評價指標是 F-measure它是精確率和召回率的調和平均能同時懲罰漏檢和誤檢?？梢栽隍炞C循環里加上這段計算precision (pred target).sum() / (pred.sum() 1e-6) recall (pred target).sum() / (target.sum() 1e-6) f_measure 2 * precision * recall / (precision recall 1e-6)需要根據驗證集上的 F-measure 來選擇后處理閾值。常見的做法是在驗證集上對每個候選閾值計算 F-measure取平均值最高的閾值作為最終閾值再對測試集做推理。這套流程雖然簡單但比固定 0.5 更貼合數據分布。6.3 從 256 到原分辨率輸出 mask 的恢復如果訓練時把原圖縮放到 256x256推理得到 256x256 的掩膜后需要把它放回原圖尺寸參與指標計算。這里有一個很容易被忽略的點cv2.resize對 mask 要用最鄰近插值但如果你用 PIL 的Image.resize默認使用 BILINEAR會產生非 0/1 的中間值。務必要顯式指定Image.NEAREST。所有掩膜后處理操作建議統一在 CPU 上用 OpenCV 完成不要用 GPU 張量直接 resize因為 PyTorch 的 interpolate 默認是雙線性輸出會變成浮點掩膜。保存結果時再強制mask[mask 0.5] 1; mask[mask 0.5] 0。這樣交付的掩膜才是純 0/1不會因為格式問題被下游指標腳本誤判。本文還有配套的精品資源點擊獲取