简介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它是精确率和召回率的调和平均能同时惩罚漏检和误检。可以在验证循环里加上这段计算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它是精确率和召回率的调和平均能同时惩罚漏检和误检。可以在验证循环里加上这段计算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不会因为格式问题被下游指标脚本误判。本文还有配套的精品资源点击获取