简介：这是一份面向图像处理入门与进阶学习者的Python实现U-Net图像分割资源，覆盖从数据准备、模型搭建、损失函数选择到训练与预测的完整流程，适合需要上手语义分割或参考现有工程代码的开发者。压缩包共21个文件，约5.6MB，包含Python脚本、README说明、多张预测结果对比图、动图演示及卫星图像分割示例，其中smooth_tiled_predictions.py可用于图像切片平滑融合，帮助改善大图分割效果。文件类型涵盖py、md、jpg、png、gif、txt与license等，目录结构清晰，便于按需查阅。目前已有10154人学习，具有较高的参考价值。通过阅读代码与示例输出，读者可理解U-Net在图像分割中的实际应用方式，并在此基础上结合自身数据集进行调整与扩展。 刚接触分割任务那阵子，我手里正好有一个户外场景需求：要把照片里的广告牌区域完整抠出来。当时我第一个想到的方案就是UNet。陆陆续续用Python把数据读取、模型搭建、训练、推理跑通之后，我越来越确信一件事：要做图像分割，UNet 是非常适合作为入门的模型，也是很多真实项目里能直接出效果的选择。
这篇文章我会把整个项目从零到一完整梳理一遍：为什么选 UNet、环境怎么配、数据集怎么准备、网络结构怎么手写、训练时有哪些坑、最后怎么把模型用在推理上。无论你是刚开始学深度学习的 Python 爱好者，还是被医学图像分割、广告牌分割这类任务折磨的研究生，这篇都能给你一份能直接上手的实操路线。
1. 项目思路拆解：用 UNet 解决分割问题的合理路径
1.1 图像分割和分类到底差在哪
很多初学者一开始容易把图像分类和图像分割混在一起。图像分类是对一整张图给出一个标签，比如“这张图里有一只猫”；而图像分割要做的就是逐像素判断类别，等于给每个像素都贴上一个标签。放到广告牌分割的场景里，就是把图片里每一块属于广告牌的区域都标成前景，其他区域标成背景。这个差异决定了模型结构很不一样：分类网络最后会接一个全局池化再加全连接层，输出一个概率分布；分割网络则需要保留空间位置信息，输出的是一张和输入尺寸相同或接近的 mask。
UNet 之所以适合分割任务，核心在于它的 U 型结构。左边是收缩路径，逐层下采样，让网络不断看到更大的感受野，学到“这是什么物体”这种高层语义；右边是扩展路径，逐层上采样，把特征图恢复到原图大小；中间用跳跃连接把左边下采样前的特征直接拼到右边对应层，这样高层的语义信息和浅层的边缘、纹理细节能同时被利用。说白了，分割任务最怕的就是边缘糊成一片，跳跃连接恰好是解决这个问题的关键设计。
1.2 跑通一个分割项目需要准备的东西
我在做这个项目时，技术栈非常简单，核心就三块：Python 做开发语言、PyTorch 做深度学习框架、OpenCV 和 Pillow 做图像处理。PyTorch 的生态足够成熟，写网络结构、写训练循环、做数据增强都有现成轮子，调试起来也比一些老框架直观得多。
完整流程大概是这样的：准备标注数据 → 数据预处理和加载 → 搭建 UNet 模型 → 定义损失函数和优化器 → 训练若干轮 → 在验证集上看效果 → 导出权重做推理。整套流程里最花时间的往往是数据准备和调参，而不是模型代码本身。环境方面，如果你有 NVIDIA 显卡，建议在官网装好 CUDA 并安装对应版本的 PyTorch；没有显卡也没关系，用 CPU 训练把图片缩到 256x256，跑几十轮同样能验证整个流程，只是速度慢一些。
2. 环境准备与数据预处理：先把地基打好
2.1 Python 环境快速配置
很多初学者容易在这里卡住，尤其是第一次配 PyTorch 的人。我的建议是直接用 Anaconda 或者 Miniconda 建一个独立环境，别看它占几个 G 磁盘，但至少不会把系统自带的 Python 环境搞坏。命令行执行：
conda create -n unet python=3.10 -y conda activate unet pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121 pip install opencv-python pillow tqdm matplotlib numpy这里的 cu121 表示 CUDA 12.1 对应的 PyTorch 版本，具体版本号要根据自己显卡驱动支持的最高 CUDA 版本确定。装完以后在 Python 里执行import torch; print(torch.cuda.is_available())，如果输出 True 就说明 GPU 可用。日常开发我习惯用 VSCode，在命令面板里选择 Python 解释器时，一定要选到刚才创建的那个 unet 环境，不然代码里 import torch 会出现“找不到模块”的问题。
注意：安装依赖包时如果遇到网络慢或者超时，可以在 pip 命令后面追加
-i https://pypi.tuna.tsinghua.edu.cn/simple换成国内镜像源，速度会快很多。
2.2 数据集的选择与 mask 的正确读法
做分割项目最难的不是模型，而是数据。自己标注数据成本太高，入门阶段我强烈建议先用公开数据集把流程跑通。我这里用的是 Oxford Pets 数据集，它包含 37 类宠物图片以及对应的前景分割标注，下载解压之后有两个重要目录：images存原图，annotations/trimaps存三分类掩膜。每个像素的取值是 1、2、3，分别表示背景、宠物主体和宠物边缘。
读取数据时最容易踩的坑是 mask 的缩放方式。原图和 mask 都要调整到统一尺寸，但resize的插值方法必须区分开：图片用双线性插值，mask 一定得用最近邻插值，否则插值会在类别交界处产生原本不存在的过渡类别，给训练带来脏标签。我常用的加载函数如下：
from PIL import Image import numpy as np import torch def load_image_and_mask(img_path, mask_path, size=(256, 256)): image = Image.open(img_path).convert("RGB").resize(size) mask = Image.open(mask_path).resize(size, resample=Image.NEAREST) image = np.array(image).astype(np.float32) / 255.0 mask = np.array(mask).astype(np.int64) # 三分类转二分类：把类别2和3合并为前景 mask = torch.from_numpy(mask) mask = (mask > 1).long() image = torch.from_numpy(image).permute(2, 0, 1) return image, mask数据加载器里还应该做归一化，很多公开教程直接用/255.0，这个对大部分自然图像数据集够用，但如果你使用的是标准化版本的预训练权重，记得要按 ImageNet 的均值和方差做标准化。训练集和验证集建议按 8:2 划分，并且固定随机种子，保证每次实验的数据分布一致，这样后面调参才有可比性。
3. UNet 核心代码实现：从零手写模型
3.1 手写一个可直接运行的 UNet 网络结构
网上 UNet 的实现版本很多，初学者最怕的就是代码太长看不下去。其实 UNet 的结构非常规律，把基础模块拆出来，整个模型就是几个模块的堆叠。核心由四部分组成：完成两次卷积的双卷积块、下采样池化、上采样转置卷积、跳跃连接拼接。我用 PyTorch 实现的原版结构如下：
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, kernel_size=3, padding=1), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True) ) def forward(self, x): return self.conv(x) class UNet(nn.Module): def __init__(self, in_channels=3, n_classes=2): super().__init__() self.enc1 = DoubleConv(in_channels, 64) self.enc2 = DoubleConv(64, 128) self.enc3 = DoubleConv(128, 256) self.enc4 = DoubleConv(256, 512) self.pool = nn.MaxPool2d(2) self.bottleneck = DoubleConv(512, 1024) self.up4 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2) self.dec4 = DoubleConv(1024, 512) self.up3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2) self.dec3 = DoubleConv(512, 256) self.up2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2) self.dec2 = DoubleConv(256, 128) self.up1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2) self.dec1 = DoubleConv(128, 64) self.out = nn.Conv2d(64, n_classes, kernel_size=1) def forward(self, x): e1 = self.enc1(x) e2 = self.enc2(self.pool(e1)) e3 = self.enc3(self.pool(e2)) e4 = self.enc4(self.pool(e3)) b = self.bottleneck(self.pool(e4)) d4 = self.dec4(torch.cat([self.up4(b), e4], dim=1)) d3 = self.dec3(torch.cat([self.up3(d4), e3], dim=1)) d2 = self.dec2(torch.cat([self.up2(d3), e2], dim=1)) d1 = self.dec1(torch.cat([self.up1(d2), e1], dim=1)) return self.out(d1)这里有几个细节值得多说一句。一是通道数从 64 开始逐层翻倍到 1024，这是 UNet 最经典的设计，实验下来在大多数据集上都稳定；二是上采样没有用简单的插值而是用转置卷积，因为转置卷积里的参数是可以学习的，能让网络自己去学会如何恢复细节；三是最后的输出层用 1x1 卷积把通道数压缩到类别数，不做激活，具体是 sigmoid 还是 softmax 由损失函数决定，这样数值上更稳定。
3.2 损失函数和评价指标选型
分割任务最常用的损失函数是交叉熵，但直接只用交叉熵会有一个明显问题：当背景占绝大多数像素时，模型很容易学成“全都预测成背景”，因为这样损失也很低。这种情况在广告牌分割和医学图像分割里都非常典型。所以我做分割项目时习惯把交叉熵和 Dice Loss 组合在一起用，Dice Loss 直接优化的是预测区域和真实区域的像素级重叠度，对正负样本不敏感。二分类场景下可以这样实现：
def dice_loss(pred, target, smooth=1e-6): pred = torch.sigmoid(pred) pred = pred.reshape(pred.size(0), -1) target = target.reshape(target.size(0), -1) intersection = (pred * target).sum(dim=1) dice = (2.0 * intersection + smooth) / (pred.sum(dim=1) + target.sum(dim=1) + smooth) return 1 - dice.mean() # 训练时组合使用 criterion = lambda pred, target: nn.CrossEntropyLoss()(pred, target) + dice_loss(pred, target.float())评价指标上，我一般看两个：Dice 系数和 IoU。IoU 的计算方式是对每个类别分别求交集除以并集，然后取平均。给一个可用的计算函数：
def compute_iou(pred_mask, true_mask, n_classes=2): ious = [] pred_mask = pred_mask.reshape(-1) true_mask = true_mask.reshape(-1) for cls in range(n_classes): pred_cls = (pred_mask == cls) true_cls = (true_mask == cls) intersection = (pred_cls & true_cls).sum().item() union = (pred_cls | true_cls).sum().item() ious.append(intersection / union if union > 0 else float("nan")) return ious这里输入的pred_mask是已经做了 argmax 的类别索引掩膜，不是概率图。很多初学者在指标计算时忘了这一步，导致计算出来的 IoU 一直不对，这也是我在群里被问过最多的一个问题。
4. 训练细节与调参实战
4.1 超参数设置和我的参数经验
训练阶段的超参数我直接给出自己常用的一套基准配置：输入尺寸 256x256，batch size 16，初始学习率 1e-4，优化器用 AdamW，训练 80 到 100 轮。显存不够时把 batch size 降到 8，同时记得顺手把学习率也降到 5e-5，否则容易震荡。固定随机种子这一步千万别省，它能让你每次复现实验都是同一套数据顺序和权重初始化，调参时才不会把参数改动的效果和随机性混在一起。
完整训练循环的骨架大致是：
for epoch in range(num_epochs): model.train() total_loss = 0.0 for images, masks in train_loader: images = images.to(device) masks = masks.to(device) outputs = model(images) loss = criterion(outputs, masks) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() val_iou = evaluate(model, val_loader, device) print(f"Epoch {epoch}: loss={total_loss / len(train_loader):.4f}, val_iou={val_iou:.4f}")在训练过程中我会观察两个曲线：训练 loss 和验证集 IoU。如果训练 loss 一直在降但验证 IoU 涨到一定程度后不动了，大概率是开始过拟合，此时可以提前停止训练，或者加大数据增强力度。如果训练 loss 一开始就下不去，那优先怀疑学习率过大导致梯度震荡，或者数据加载环节出了问题，比如 mask 和图片没有对齐。
4.2 数据不平衡、过拟合、学习率失控三个坑
先说数据不平衡。广告牌分割里背景占面积往往超过 90%，训练初期模型最可能的行为就是全预测成背景。我的处理思路有两层：一是损失函数里引入 Dice Loss，让前景区域的重要性和背景平分秋色；二是在数据层面做类别采样，让包含前景的样本在训练集中占更高比例。实测下来，组合损失比单一交叉熵在广告牌这类任务上提升非常明显。
过拟合是另一个高频问题。分割任务的模型参数量不小，公开数据集通常又只有几千张图，训练轮数一多，验证集 IoU 反而掉头向下。我的经验是：先用随机水平翻转、随机旋转 10 度、随机亮度对比度扰动这三件套做数据增强，如果还是过拟合，再加 CutOut 或者随机遮挡。尽量少用 Dropout，UNet 在卷积层里加 Dropout 对分割效果的帮助并不明显，反而拖慢收敛。
学习率方面，我一般搭配ReduceLROnPlateau调度器，当验证 IoU 连续 5 个 epoch 不涨时自动把学习率乘以 0.5。这样前期用稍大的学习率快速下降，后期自动小步慢走找最优解，比固定学习率省心得多。
5. 推理与后处理：让输出结果能用起来
5.1 推理流程与预处理一致性
训练完成以后，模型要真正用于推理，最容易出错的就是“预处理不一致”。我在第一次部署时就栽过这个跟头：训练时把图片 resize 到 256x256 再归一化，推理时却直接用原始尺寸塞进模型，结果输出 mask 尺寸对不上，程序直接报错。推理阶段有一套固定写法，先加载权重，再把模型切到 eval 模式，在torch.no_grad()下让梯度不参与计算，这样又快又省显存：
model.load_state_dict(torch.load("unet_best.pth", map_location="cpu")) model.to(device).eval() with torch.no_grad(): image_tensor = load_image_and_mask(...)[0].unsqueeze(0).to(device) logits = model(image_tensor) pred_mask = torch.softmax(logits, dim=1).argmax(dim=1).squeeze(0).cpu().numpy()注意model.eval()非常关键，如果不写，BatchNorm 层仍然会使用训练阶段统计的均值方差去计算，结果就是在验证集上看着挺好、实际推理时输出完全不可用。这类问题藏得很深，很多人排查半天都找不到原因。
5.2 后处理技巧：抠图更干净
模型输出的原始 mask 通常会有一些细碎的噪点，直接叠到原图上勉强能看，但如果要用于后续业务处理，比如统计广告牌像素占比或计算中心点坐标，就要做后处理。我常用的两步是连通域过滤和形态学操作。先用 OpenCV 的connectedComponentsWithStats把面积小于一定阈值的连通域当成噪点去掉，再用腐蚀膨胀修一下边缘毛刺。针对广告牌这种块状目标，拉直边缘效果尤其明显：
import cv2 import numpy as np num, labels, stats, centroids = cv2.connectedComponentsWithStats(pred_mask, connectivity=8) clean_mask = np.zeros_like(pred_mask) min_area = 100 for i in range(1, num): if stats[i, cv2.CC_STAT_AREA] >= min_area: clean_mask[labels == i] = 1 clean_mask = cv2.morphologyEx(clean_mask, cv2.MORPH_CLOSE, np.ones((5, 5), np.uint8))如果做的是多类别分割，后处理前记得把每个类别拆出来单独处理，最后再按类别编号合并，否则连通域统计会把不同类别粘连在一起，导致过滤逻辑失效。
6. 常见问题与排错速查
| 现象 | 可能原因 | 解决建议 |
|---|---|---|
import torch报 ModuleNotFoundError | 当前解释器不在 conda 环境里 | 在 VSCode 右下角切换解释器，或命令行先conda activate unet |
torch.cuda.is_available()返回 False | 安装的 PyTorch 是 CPU 版或 CUDA 不匹配 | 删除后按显卡驱动支持的 CUDA 版本重装对应包 |
| 训练时显存不足 OOM | batch size 太大或输入图过大 | 减小 batch size、降低输入分辨率，或用 gradient accumulation |
| 训练到一半 loss 变成 NaN | 学习率过大、数据标签有异常值、损失函数写错 | 调低学习率，检查 mask 取值是否超出类别范围 |
| 模型输出全黑或全白 | 推理前忘了model.eval()，或 softmax 后没做 argmax，或预处理不一致 | 检查模型模式、预测后处理和数据归一化链路 |
| Windows 下 DataLoader 报多进程错误 | num_workers在 Windows 下需要特殊处理 | 将num_workers设为 0，或把训练代码放到if __name__ == "__main__":里 |
| pip 安装包下载慢或超时 | 网络原因 | 使用国内镜像源-i https://pypi.tuna.tsinghua.edu.cn/simple |
| 验证集准确率很高但 IoU 很低 | 类别严重不平衡，模型学成了全背景 | 引入 Dice Loss 或调整权重 |
这套排查表基本覆盖了我在实际项目中遇到的大部分低级错误。很多问题不是模型设计上的问题，而是环境或数据细节没处理好，所以我建议第一次跑通项目后不要急着改结构，先把整条链路的每个环节都验证一遍，后续做任何改动心里都有底。
7. 除了入门分割，UNet 还能往哪走
7.1 医学图像分割为什么爱用 UNet
UNet 最初就是为医学图像分割设计的，也因此在这个领域影响最深。医学影像数据标注成本极高，往往一张像样的标注图需要医生手工描半天，样本量经常只有几十到几百张，这种情况下那些只能在大型数据集上发挥威力的大模型根本跑不起来。UNet 因为结构紧凑、参数适中，加上跳跃连接能充分利用有限的标注信息，在少量数据上也能学到一个可用的分割结果。如果你之后要做器官分割、病灶检测，UNet 及其变体依然是很可靠的首选基线。
7.2 几个好上手的改进方向
等基础版本跑通以后，你可以往两个方向做改进。第一个方向是替换编码器，把 UNet 左边的卷积块换成 MobileNet 或 EfficientNet 这类轻量骨干网络，先用 ImageNet 预训练权重初始化，再在自己的任务上微调，通常能带来明显的 IoU 提升，训练时间也会缩短。第二个方向是加注意力机制，比如在跳跃连接前加一个 SE 模块，或者直接参考 Attention U-Net 的做法，让网络在拼接特征时自动抑制无关区域、突出关键区域。改模型结构时有一个原则：每次只改一个模块，跑一组实验验证效果，不要一口气同时换骨干、加注意力、换损失函数，否则效果变好了你都不知道是哪个改动起了作用。
7.3 工程化部署时还需要考虑什么
如果你的目标是把训练好的模型推到实际系统里，那还需要考虑速度。PyTorch 模型直接部署存在依赖重、启动慢的问题，常见做法是先把模型导出成 ONNX，再用 ONNX Runtime 或 TensorRT 做推理加速。导出时要注意固定输入尺寸和动态轴，如果输入图片分辨率会变化，需要把torch.onnx.export中的 dynamic_axes 参数设置好。另外，如果模型本身就是几 MB 的小模型，CPU 上的 ONNX Runtime 推理速度已经足够满足很多实时场景，不一定非得买一块高端显卡。
做分割项目这段时间，我最大的体会就是：千万别被各种改进技巧迷了眼。先老老实实把最基础的 UNet 跑通，亲眼看一遍 IoU 曲线是怎么从 0.2 涨到 0.8 的，再去研究注意力、多尺度融合这些改进点，你会突然发现原来那些论文里写的结构，本质上都是在解决你自己踩过的那些坑。最后再分享一个很实用的小习惯：每次训练把超参数、数据集划分方式、loss 曲线截图都记录成日志，同一个模型同一个数据集，改了学习率就可能差出好几个点，没有记录的话，回溯起来真的非常痛苦。
本文还有配套的精品资源，点击获取
