UNet图像分割实战:从代码拆解到训练推理与二次开发

发布时间:2026/10/24 2:34:42
UNet图像分割实战:从代码拆解到训练推理与二次开发 简介这份资源面向深度学习初学者与计算机视觉方向的学习者围绕经典 Unet 卷积神经网络完成图像分割任务可用于医学影像分析、遥感图像处理等场景的入门实践与课程复现。压缩包共 253 个文件约 32.34MB以 243 个 png 图像样本为主体另含 3 个 py 脚本、3 个 tif 影像、2 个 ipynb 交互式笔记本及说明文档覆盖数据准备、模型搭建与训练流程。已有 4339 人学习下载具备一定的参考基础。读者可从中获得 Unet 收缩路径与扩张路径的完整实现思路包括跳跃连接的特征融合方式、交叉熵损失与优化器配置、数据预处理与增强策略以及训练验证流程和分割结果可视化配套笔记本便于逐步运行调试图像样本可直接用于复现实验帮助理解上下文信息与细节保留之间的平衡并为进一步引入多尺度融合或注意力机制打下基础。1. 拿到一个 unet 图像分割包先别急着跑上周有个做医学影像的朋友丢过来一个压缩包说里面是 unet 网络实现图像分割的完整代码让我帮忙看看能不能跑。我打开一看目录里躺着模型定义、训练脚本、数据加载、推理入口典型的深度学习工程结构。这类资源对想入门图像分割的开发者来说价值不小——它把 unet 这个经典网络从论文里的结构图落成了能训练、能推理、能改的工程代码。适合谁刚接触深度学习想找个完整项目练手的、做医学影像或遥感分割需要 baseline 的、以及想理解编码器-解码器架构怎么在代码里落地的开发者。这篇就按我拆包的顺序从项目定位、目录结构、核心模块到怎么跑、怎么改一层层讲清楚。2. 项目定位与技术选型unet 为什么还是分割任务的常青树2.1 这个包解决的是什么问题图像分割的本质是给每个像素分类。unet 2015 年提出时是为医学影像分割设计的结构上是标准的编码器-解码器编码器一路下采样提特征解码器一路上采样恢复分辨率中间用跳跃连接把浅层的高分辨率特征拼回深层语义特征。这个设计的好处是小目标、边缘细节不会在下采样过程中丢干净对细胞、病灶、道路这类细长或小面积目标特别友好。这个资源包把上述结构完整实现了一遍包含数据读取、模型定义、损失函数、训练循环、验证指标、推理可视化几个部分。它不是某个框架的官方示例而是一个可以独立跑起来的工程拿到手改改数据路径就能训练自己的数据集。对想理解「一个分割项目到底由哪些文件组成」的人来说这种完整度比看十篇论文都管用。2.2 技术栈与依赖判断从目录里的文件命名和导入语句看这个包基于 PyTorch 实现配套 torchvision 做数据变换numpy 和 PIL 处理图像读写matplotlib 做训练曲线和分割结果可视化。这是目前分割项目最主流的组合生态成熟、调试方便。常见做法是再引入 albumentations 做数据增强但如果包里没带用 torchvision.transforms 手写几个也够用。依赖大致如下我一般会先建虚拟环境再装# 建议 Python 3.8 以上PyTorch 按自己显卡的 CUDA 版本装 pip install torch torchvision pip install numpy pillow matplotlib # 如果代码里用了增强库再补一条 pip install albumentations这里有个血泪经验PyTorch 版本和 CUDA 驱动不匹配是新手翻车率最高的地方。装之前先nvidia-smi看驱动支持的 CUDA 上限再去 PyTorch 官网找对应版本的安装命令别直接pip install torch了事否则训练时一句CUDA error能卡你半天。2.3 和同类方案的差别有人会问现在分割模型那么多deeplab、segformer、sam 都出来了为什么还看 unet我的判断是unet 结构简单、参数量小、训练成本低在小数据集上往往比大模型更稳。医学影像动辄几百张图上大模型容易过拟合unet 反而是性价比最高的 baseline。这个包的价值不在于刷榜而在于给你一个干净、可读、可改的起点——你先把它跑通理解每个模块在干什么再去换 backbone、加注意力、改损失函数路径是清晰的。3. 目录结构与核心模块拆解一个分割工程由哪些文件组成3.1 典型目录长什么样解压后我看到的目录结构大致是这样不同版本可能略有出入但骨架一致unet_segmentation/ ├── datasets/ # 数据读取与预处理 │ └── dataset.py ├── models/ # 网络结构定义 │ └── unet.py ├── losses/ # 损失函数 │ └── dice_loss.py ├── utils/ # 指标、可视化、日志 │ ├── metrics.py │ └── visualize.py ├── train.py # 训练入口 ├── predict.py # 推理入口 ├── config.py # 超参与路径配置 └── requirements.txt这个划分方式很工程化数据、模型、损失、工具各自独立入口脚本只负责串联。二次开发时你改哪块就动哪个文件不会牵一发动全身。我特别看重config.py这个设计把学习率、batch size、图像尺寸、数据路径集中管理比散落在各个脚本里的硬编码强太多。3.2 模型定义编码器与解码器怎么拼unet 的核心在models/unet.py。下面是我从包里摘出来的结构做了注释精简import torch import torch.nn as nn class DoubleConv(nn.Module): 两次 3x3 卷积 BN ReLUunet 的基本单元 def __init__(self, in_ch, out_ch): super().__init__() self.block 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.block(x) class UNet(nn.Module): def __init__(self, in_channels3, num_classes1): super().__init__() # 编码器每次先卷积再池化通道翻倍 self.down1 DoubleConv(in_channels, 64) self.down2 DoubleConv(64, 128) self.down3 DoubleConv(128, 256) self.down4 DoubleConv(256, 512) self.pool nn.MaxPool2d(2) # 解码器上采样后与编码器特征拼接 self.up1 nn.ConvTranspose2d(512, 256, 2, stride2) self.conv1 DoubleConv(512, 256) self.up2 nn.ConvTranspose2d(256, 128, 2, stride2) self.conv2 DoubleConv(256, 128) self.up3 nn.ConvTranspose2d(128, 64, 2, stride2) self.conv3 DoubleConv(128, 64) # 输出层1x1 卷积映射到类别数 self.out nn.Conv2d(64, num_classes, 1) def forward(self, x): # 编码路径保存每一级特征用于跳跃连接 d1 self.down1(x) d2 self.down2(self.pool(d1)) d3 self.down3(self.pool(d2)) d4 self.down4(self.pool(d3)) # 解码路径上采样后与对应层级特征 concat u1 self.up1(d4) u1 self.conv1(torch.cat([u1, d3], dim1)) u2 self.up2(u1) u2 self.conv2(torch.cat([u2, d2], dim1)) u3 self.up3(u2) u3 self.conv3(torch.cat([u3, d1], dim1)) return self.out(u3)逻辑说明DoubleConv是重复使用的卷积块保证每次特征提取都是两次卷积加归一化。编码器每经过一个DoubleConv就池化一次分辨率减半、通道翻倍。解码器用转置卷积上采样torch.cat把上采样结果和编码器同层特征在通道维拼接这就是跳跃连接在代码里的样子。参数上in_channels按输入图像通道改RGB 是 3灰度是 1num_classes二分类填 1多分类填类别数。注意二分类时输出层通常配 sigmoid多分类配 softmax这个在损失函数里体现。3.3 数据加载与增强datasets/dataset.py负责把图像和掩码配对读进来。分割任务的数据集有个特点输入是图像标签是同尺寸的掩码两者必须做完全一致的几何变换否则图像翻转了掩码没翻训练直接崩。常见做法是自定义 Dataset 类from torch.utils.data import Dataset from PIL import Image import torchvision.transforms.functional as TF import random class SegDataset(Dataset): def __init__(self, img_paths, mask_paths, size256): self.img_paths img_paths self.mask_paths mask_paths self.size size def __len__(self): return len(self.img_paths) def __getitem__(self, idx): img Image.open(self.img_paths[idx]).convert(RGB) mask Image.open(self.mask_paths[idx]).convert(L) img TF.resize(img, [self.size, self.size]) mask TF.resize(mask, [self.size, self.size]) # 同步随机翻转图像和掩码必须一起动 if random.random() 0.5: img TF.hflip(img) mask TF.hflip(mask) img TF.to_tensor(img) mask TF.to_tensor(mask) mask (mask 0.5).float() # 二值化掩码 return img, mask这里的关键点是几何增强必须同步。我见过有人用transforms.Compose分别处理图像和掩码结果随机裁剪的坐标不一致训练 loss 一直不降排查半天才发现是数据对不上。参数size控制输入分辨率unet 要求输入尺寸能被 16 整除四次下采样256 或 512 都是安全值。3.4 损失函数与评估指标分割任务常用 Dice Loss 或交叉熵与 Dice 的组合。包里losses/dice_loss.py的实现大致是import torch import torch.nn as nn class DiceLoss(nn.Module): def __init__(self, smooth1.0): super().__init__() self.smooth smooth def forward(self, pred, target): pred torch.sigmoid(pred) # 二分类输出转概率 pred pred.view(-1) target target.view(-1) intersection (pred * target).sum() dice (2 * intersection self.smooth) / ( pred.sum() target.sum() self.smooth ) return 1 - dicesmooth是平滑项防止分母为零。Dice 系数衡量预测和真值的重叠度越接近 1 越好所以损失取1 - dice。实际训练时我一般用0.5 * BCE 0.5 * DiceBCE 提供稳定的梯度Dice 直接优化重叠指标两者互补。评估指标在utils/metrics.py里通常是 IoU 和 Dice验证时算一下就知道模型有没有在学。4. 从零跑通训练、验证与推理的完整流程4.1 数据准备与配置修改跑之前先把数据按包里的约定组织好。常见约定是图像放一个目录、掩码放另一个目录文件名一一对应。然后打开config.py改这几项# config.py 示例 DATA_ROOT ./data # 数据集根目录 IMG_DIR images # 图像子目录 MASK_DIR masks # 掩码子目录 IMG_SIZE 256 # 输入分辨率需被 16 整除 BATCH_SIZE 8 # 显存不够就调小 LR 1e-4 # 学习率unet 常用 1e-4 到 1e-3 EPOCHS 50 # 训练轮数 NUM_CLASSES 1 # 二分类为 1 DEVICE cuda # 没显卡改 cpu改完先跑一个 epoch 看看 loss 有没有下降。如果 loss 是 nan八成是学习率太大或者掩码没归一化如果 loss 不动检查数据和标签是不是对错了位置。4.2 训练脚本的执行逻辑train.py的主循环结构是标准的遍历 epoch遍历 dataloader前向、算损失、反向、更新每个 epoch 结束在验证集上算指标、存最优模型。执行命令python train.py --config config.py训练过程中重点看两个信号训练 loss 是否稳定下降验证 Dice 是否同步上升。如果训练 loss 降但验证指标不涨说明过拟合了加数据增强、减模型容量或者早停。如果两者都不动先怀疑数据管道。我一般会在第一个 epoch 结束后打印几张预测掩码和真值的对比图肉眼确认模型输出是不是全黑或全白这比盯数字直观。4.3 推理与结果可视化训练完用predict.py对新图像做推理python predict.py --checkpoint best.pth --input ./test_images --output ./results推理脚本做的事是加载模型权重、切到 eval 模式、对输入做和训练一致的预处理、前向得到 logits、sigmoid 后按阈值二值化、保存掩码或叠加可视化。注意 eval 模式和torch.no_grad()一定要加否则 BN 层用 batch 统计量、显存还白白占用。阈值默认 0.5如果发现小目标漏检多可以降到 0.3 试试这个参数在医学影像里经常需要按场景调。5. 避坑与常见问题跑 unet 时最容易翻车的几个地方5.1 输入尺寸不能被 16 整除现象前向传播时报张量维度不匹配torch.cat拼接失败。原因unet 有四次下采样每次分辨率减半如果输入尺寸不是 16 的倍数上采样回来尺寸和编码器特征对不上。解决把IMG_SIZE设成 256、512 这类 16 的倍数或者在模型里加 padding 保证整除。5.2 掩码没二值化导致 loss 异常现象Dice loss 一直在 0.5 附近晃模型输出模糊。原因掩码读进来是 0-255 的灰度图没归一化到 0/1和 sigmoid 输出不在一个量纲。解决在 Dataset 里加mask (mask 0.5).float()确保标签是严格的 0 和 1。5.3 训练和推理预处理不一致现象训练指标很好推理结果一塌糊涂。原因训练时做了归一化比如除以 255 再减均值推理脚本忘了同样的处理。解决把预处理逻辑抽成一个函数训练和推理共用别各写各的。5.4 显存不够直接改 batch size 之外还要看分辨率现象CUDA out of memory。原因显存占用和 batch size、图像分辨率、模型通道数都相关。解决优先降 batch size 到 1 试试还不行就降分辨率再不行考虑混合精度训练。别一上来就换显卡很多时候调参就能解决。5.5 类别不平衡导致全预测背景现象Dice 看着不低但可视化出来全是黑的。原因背景像素远多于前景模型学会了偷懒全预测背景。解决用 Dice loss 替代纯交叉熵或者在交叉熵里加类别权重再配合重采样。6. 二次开发换 backbone、改损失、加注意力的实操思路跑通之后这个包真正的价值在于它是一个可改造的骨架。我一般按三个方向动刀按改动成本从低到高排。第一个方向是换损失函数。把losses/下的 Dice 和 BCE 组合起来权重从 0.5/0.5 开始调。如果数据里小目标多可以试试 Focal Loss 加 DiceFocal 压制易分样本的梯度让模型聚焦难样本。改的时候只动损失文件训练脚本里换一行实例化就行风险最低。第二个方向是换编码器。原版 unet 的编码器是手写卷积可以换成预训练的 ResNet 或 EfficientNet 做 backbone解码器保留。这样在小数据集上能借预训练权重的力收敛快很多。改法是新建一个模型文件用 torchvision 加载 backbone取中间层输出接回原来的解码器注意通道数要对齐。下面是个示意import torchvision.models as models import torch.nn as nn class ResNetUNet(nn.Module): def __init__(self, num_classes1): super().__init__() backbone models.resnet34(pretrainedTrue) # 取前几层作为编码器具体层名按 torchvision 版本确认 self.encoder nn.Sequential(*list(backbone.children())[:8]) # 解码器部分沿用原 unet 的上采样结构 # ... 此处接 DoubleConv 和转置卷积 def forward(self, x): features self.encoder(x) # 接解码器 return features第三个方向是加注意力模块。在跳跃连接拼接之后加一个 SE 块或 CBAM让模型自己学哪些通道重要。改动点集中在DoubleConv后面插入几行代码即可但要注意加了模块后参数量上升小数据集上可能反而过拟合加完一定要对比验证指标。验证改造是否有效我的习惯是固定随机种子、固定数据划分只改一个变量跑三次取平均。别一次改好几个地方否则指标涨了都不知道是谁的功劳。从那以后我每次改模型结构都强制先跑一遍原始 baseline 存下指标再动代码不然改着改着就忘了起点在哪。希望这个拆解能帮你把这个 unet 包真正跑起来、改起来。本文还有配套的精品资源点击获取

关于本文作者

来自尧图内容编辑团队

尧图内容编辑团队 内容团队

尧图内容编辑团队

本文由尧图网络内容编辑团队执笔。团队由资深项目经理、前端工程师与设计师组成,所有内容均来自亲手交付的真实项目,先讲清问题、再给出可落地的解法。尧图深耕北京网站建设十年,服务过京华建材集团、智造科技等各行业客户,把一线经验沉淀为可复用的行业观察。

  • 十年建站经验,覆盖建材、制造、服务、文创等
  • 项目经理把关选题与事实准确性
  • 工程师与设计师联合撰写专业细节
  • 统一编辑规范,保证文风与排版一致
  • 每月复盘转化数据,迭代选题方向

延伸阅读

相关资讯与近期热门内容

深度阅读推荐

建站决策前值得细读的三篇

网站改版的5个关键决策
2024-08-12

网站改版的5个关键决策

什么时候该改版、改到什么程度、如何避免流量掉光,京华建材集团改版复盘给出答案。

获取专属建站方案

看完文章,把您的行业与预算告诉我们,免费获取一份量身定制的官网建设方案与报价。

立即免费咨询