简介：这份资源面向深度学习初学者与计算机视觉方向的开发者，提供一套可运行的UNet图像分割完整实现，帮助读者理解编码器—解码器结构与跳跃连接如何兼顾上下文信息与细节保留，适用于医学影像、遥感等像素级分类场景的学习与复现。压缩包共253个文件，约32.34MB，以243张png图像数据为主体，辅以3个py脚本、2个ipynb交互式笔记本、3个tif影像及说明文档，覆盖数据准备、模型搭建、训练与预测全流程。目前已有4339人学习下载，具备较好的参考基础。读者可从中获得网络结构定义、损失函数与优化器配置、数据增强策略，以及训练验证流程的完整代码，并借助可视化分割结果与Dice、IoU等评估指标对照实验效果，快速搭建自己的分割实验环境并在此基础上调整超参数或引入注意力机制等改进。
1. 从一张标注图说起：这个 U-Net 项目到底能干什么
手里有一批医学影像或者遥感图，想自动把目标区域抠出来，第一反应往往是找现成的分割模型。这个压缩包给的就是一套能直接跑的 U-Net 图像分割实现，输入原图、输出同尺寸的掩码，训练、推理、评估的链路都在里面。它适合两类人：一类是刚接触深度学习图像分割、想找个结构干净的项目把流程走通的新手；另一类是手里已有业务数据、想拿它当基线快速验证效果再二次开发的熟手。U-Net 本身是编码器-解码器加跳跃连接的经典结构，参数量不大、对小样本友好，这也是它在分割任务里一直被当作首选基线的原因。下面按「能干什么、怎么跑、怎么改」的顺序拆开讲。
2. 项目结构与技术选型：从数据目录到模型输出的完整链路
2.1 目录组织与各模块职责
拿到压缩包解压后，先别急着敲命令，花两分钟把目录看一遍，后面改代码会省很多事。这类 U-Net 项目通常按「数据、模型、训练、推理、工具」五块切分，结构大致如下：
unet-segmentation/ ├── data/ │ ├── train/ │ │ ├── images/ # 训练原图 │ │ └── masks/ # 对应标注掩码，文件名与原图一一对应 │ └── val/ │ ├── images/ │ └── masks/ ├── models/ │ └── unet.py # U-Net 网络定义 ├── datasets/ │ └── segmentation.py # Dataset 与预处理 ├── train.py # 训练入口 ├── predict.py # 单图/批量推理 ├── evaluate.py # 指标评估 ├── utils/ │ ├── metrics.py # Dice、IoU 等 │ └── visualize.py # 结果可视化 ├── configs/ │ └── default.yaml # 超参与路径配置 └── requirements.txt这个划分的好处是职责清晰：models/只管网络结构，datasets/只管数据读取和增强，train.py负责把两者串起来。二次开发时想换骨干网络只动models/，想加数据增强只动datasets/，互不干扰。常见做法是把所有路径和超参收进configs/default.yaml，避免在代码里硬编码，这一点在多人协作时尤其重要。
2.2 技术栈与依赖
核心依赖是 PyTorch + torchvision，图像读写用 OpenCV 或 PIL，指标计算和可视化用 NumPy、Matplotlib。requirements.txt里一般会锁定版本，比如：
torch>=1.10 torchvision>=0.11 numpy>=1.21 opencv-python>=4.5 matplotlib>=3.4 pyyaml>=5.4 tqdm>=4.62选 PyTorch 而不是别的框架，主要是调试友好——训练循环里想打印什么、想改哪一步都直观，对新手排查问题友好。如果你的环境已经有 CUDA，装对应版本的 torch 即可；没有 GPU 也能跑，只是训练慢，小分辨率数据在 CPU 上验证流程完全够用。
2.3 U-Net 网络结构拆解
U-Net 的核心是「下采样提特征、上采样恢复分辨率、跳跃连接补细节」。编码器每下采样一次通道翻倍，解码器每上采样一次通道减半，同时把编码器同层的特征拼过来，保证边缘细节不丢。下面是models/unet.py里最典型的实现骨架：
import torch import torch.nn as nn class DoubleConv(nn.Module): """两次 3x3 卷积 + BN + ReLU，U-Net 的基本单元""" def __init__(self, in_ch, out_ch): super().__init__() self.block = nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), ) def forward(self, x): return self.block(x) class UNet(nn.Module): def __init__(self, in_channels=3, num_classes=1, base_ch=64): super().__init__() # 编码器：4 次下采样 self.enc1 = DoubleConv(in_channels, base_ch) self.enc2 = DoubleConv(base_ch, base_ch * 2) self.enc3 = DoubleConv(base_ch * 2, base_ch * 4) self.enc4 = DoubleConv(base_ch * 4, base_ch * 8) self.pool = nn.MaxPool2d(2) # 瓶颈层 self.bottleneck = DoubleConv(base_ch * 8, base_ch * 16) # 解码器：转置卷积上采样 + 跳跃连接拼接 self.up4 = nn.ConvTranspose2d(base_ch * 16, base_ch * 8, 2, stride=2) self.dec4 = DoubleConv(base_ch * 16, base_ch * 8) self.up3 = nn.ConvTranspose2d(base_ch * 8, base_ch * 4, 2, stride=2) self.dec3 = DoubleConv(base_ch * 8, base_ch * 4) self.up2 = nn.ConvTranspose2d(base_ch * 4, base_ch * 2, 2, stride=2) self.dec2 = DoubleConv(base_ch * 4, base_ch * 2) self.up1 = nn.ConvTranspose2d(base_ch * 2, base_ch, 2, stride=2) self.dec1 = DoubleConv(base_ch * 2, base_ch) # 1x1 卷积输出每个像素的类别 logits self.out = nn.Conv2d(base_ch, num_classes, 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)几个参数要留意：in_channels按输入图通道数改，灰度图填 1、RGB 填 3；num_classes是分割类别数，二分类填 1 配合 Sigmoid，多分类填类别数配合 Softmax；base_ch控制整体宽度，显存不够就往下调，64 是常见起点。torch.cat的dim=1是通道维拼接，这是跳跃连接的关键，拼接后通道数翻倍，所以后面的DoubleConv输入通道要按翻倍后的值写，写错了会直接报维度不匹配。
2.4 数据加载与预处理
数据这块最容易翻车。掩码必须是单通道、像素值是类别索引（0/1/2…），不能是彩色图直接读进来。datasets/segmentation.py里通常这样写：
import os import cv2 import torch from torch.utils.data import Dataset class SegDataset(Dataset): def __init__(self, img_dir, mask_dir, transform=None): self.img_dir = img_dir self.mask_dir = mask_dir self.names = sorted(os.listdir(img_dir)) # 保证图与掩码顺序一致 self.transform = transform def __len__(self): return len(self.names) def __getitem__(self, idx): name = self.names[idx] img = cv2.imread(os.path.join(self.img_dir, name)) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 统一成 RGB mask = cv2.imread(os.path.join(self.mask_dir, name), cv2.IMREAD_GRAYSCALE) mask = (mask > 127).astype("float32") # 二值化，按需改阈值 if self.transform: augmented = self.transform(image=img, mask=mask) # 图与掩码同步增强 img, mask = augmented["image"], augmented["mask"] img = torch.from_numpy(img).permute(2, 0, 1).float() / 255.0 mask = torch.from_numpy(mask).unsqueeze(0) # 增加通道维 return img, mask这里有两个细节：一是sorted保证原图和掩码按同一顺序取，文件名对不上是最常见的「训练 loss 不降」原因；二是增强必须对图和掩码同步做，旋转、翻转、缩放都要用同一组随机参数，否则掩码和图像错位，模型学到的全是噪声。用 albumentations 这类库时，把 image 和 mask 一起传进去就能保证同步。
3. 跑起来：训练、推理与评估的完整操作
3.1 环境准备与数据摆放
先把依赖装上，建议用虚拟环境隔离：
python -m venv venv source venv/bin/activate # Windows 用 venv\Scripts\activate pip install -r requirements.txt数据按第 2 章的目录结构摆好，data/train/images和data/train/masks里的文件名必须完全一致，扩展名也要一致。如果原始数据是彩色掩码，先写个小脚本转成灰度索引图再放进去，别指望训练代码帮你转。数据量小的时候（比如几百张），建议按 8:2 切训练和验证，验证集单独放，不要混进训练。
3.2 训练脚本与关键参数
train.py的主循环大致长这样，重点是损失函数和优化器的选择：
import torch import torch.nn as nn from torch.utils.data import DataLoader from models.unet import UNet from datasets.segmentation import SegDataset device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = UNet(in_channels=3, num_classes=1, base_ch=64).to(device) train_ds = SegDataset("data/train/images", "data/train/masks", transform=train_aug) train_loader = DataLoader(train_ds, batch_size=8, shuffle=True, num_workers=4) # 二分类用 BCEWithLogitsLoss，多分类换 CrossEntropyLoss criterion = nn.BCEWithLogitsLoss() optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) for epoch in range(50): model.train() total_loss = 0.0 for img, mask in train_loader: img, mask = img.to(device), mask.to(device) optimizer.zero_grad() logits = model(img) loss = criterion(logits, mask) loss.backward() optimizer.step() total_loss += loss.item() print(f"epoch {epoch}, loss {total_loss / len(train_loader):.4f}") torch.save(model.state_dict(), "checkpoints/last.pth")batch_size受显存限制，8 或 16 是常见起点，显存爆了就减半；lr=1e-3配 Adam 是稳妥选择，训练后期 loss 震荡可以换ReduceLROnPlateau自动降学习率；num_workers在 Windows 上设 0 更稳，Linux 上设 4 左右。损失函数方面，二分类用BCEWithLogitsLoss（内部带 Sigmoid，别在外面再加一层），类别极不平衡时换成 Dice Loss 或 BCE+Dice 组合，效果通常更稳。
3.3 推理与结果保存
训练完拿predict.py跑单图或批量推理，核心是把 logits 过 Sigmoid 再二值化：
import cv2 import numpy as np import torch model.load_state_dict(torch.load("checkpoints/last.pth", map_location=device)) model.eval() with torch.no_grad(): img = cv2.imread("test.png") img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) tensor = torch.from_numpy(img_rgb).permute(2, 0, 1).float().unsqueeze(0) / 255.0 logits = model(tensor.to(device)) prob = torch.sigmoid(logits).squeeze().cpu().numpy() pred = (prob > 0.5).astype(np.uint8) * 255 # 阈值 0.5 可按验证集调 cv2.imwrite("pred_mask.png", pred)阈值 0.5 不是铁律，验证集上扫一遍 0.3～0.7，挑 Dice 最高的那个存进配置。推理时记得model.eval()和torch.no_grad()，前者影响 BN 层行为，后者省显存，漏了任何一个结果都可能对不上训练时的表现。
3.4 评估指标怎么看
evaluate.py一般算 Dice 和 IoU 两个指标。Dice 对前景占比敏感，IoU 更严格，两个一起看能判断模型是「漏检多」还是「误检多」。验证集 Dice 卡在 0.6 上不去，先查掩码有没有错位、类别是否平衡，再考虑换损失函数或加增强，别一上来就调网络结构。
4. 避坑与排查：训练不收敛、掩码错位这些坑
4.1 现象：loss 一直不降或剧烈震荡
原因通常是学习率过大、掩码和图像没对齐、或者标签值不在预期范围。解决顺序是：先打印一个 batch 的图和掩码肉眼核对，确认掩码是单通道且值只有 0/1；再把 lr 降到 1e-4 试几个 epoch；最后检查BCEWithLogitsLoss外面是不是多套了一层 Sigmoid，多套会导致梯度异常。
4.2 现象：训练集 Dice 很高，验证集一塌糊涂
典型的过拟合。数据量小的时候尤其明显。解决上先加数据增强（随机翻转、旋转、亮度扰动），再考虑加 Dropout 或权重衰减，最后才是减小base_ch。别急着上复杂结构，U-Net 在小数据上过拟合是常态，增强比换模型管用。
4.3 现象：推理结果全是 0 或全是 1
多半是归一化不一致。训练时图像除了 255，推理时忘了除，或者反过来。还有一种情况是掩码在训练时被二值化成了 0/1，推理输出却按 0/255 保存，可视化看着全黑。统一在 Dataset 和 predict 里用同一套预处理，把归一化参数写进配置，两边都读配置，能避免这类玄学问题。
4.4 现象：显存不够，batch_size 降到 1 还爆
检查输入分辨率是不是太大，U-Net 对分辨率敏感，1024×1024 起步就很吃显存。解决上可以先把图 resize 到 512 或 256 训练，推理时再上采样回原尺寸；或者用混合精度训练（torch.cuda.amp），显存能省近一半。另外num_workers开太大也会占内存，适当调小。
4.5 现象：多分类任务输出通道对不上
num_classes填的是类别数，但很多人忘了背景类也算一类。三分类（背景+两类目标）要填 3，配合CrossEntropyLoss，标签是 0/1/2 的整型张量，不是 one-hot。填成 2 会直接报维度错误或者静默算错，这个坑我见过不止一次。
5. 二次开发：换骨干、加损失、做部署的改造思路
拿到这套代码，真正有价值的是在它基础上改。第一个改造点是换骨干网络：把编码器换成 ResNet 或 EfficientNet 的预训练权重，小数据上收敛会快很多，改法是保留enc1~enc4的接口，内部换成对应 stage 的输出，注意记录每层输出通道数，解码器的拼接维度要跟着改。第二个改造点是损失函数，类别不平衡时把BCEWithLogitsLoss换成0.5 * BCE + 0.5 * Dice，Dice 部分自己实现十几行就够，验证集上对比一下再决定留哪个。
第三个改造点是导出部署。训练完的.pth只是权重，线上服务一般要转 ONNX：
dummy = torch.randn(1, 3, 512, 512).to(device) torch.onnx.export( model, dummy, "unet.onnx", input_names=["input"], output_names=["logits"], dynamic_axes={"input": {0: "batch"}, "logits": {0: "batch"}}, opset_version=11, )导出后拿 onnxruntime 跑一遍，和 PyTorch 输出逐像素比对，误差在 1e-4 以内才算过。动态 batch 维一定要开，否则线上只能单张推理。改完任何一处，我的习惯是固定一个 20 张的小验证集，每次改动都跑一遍 Dice 和可视化，确认没退化再继续，不然改着改着就不知道哪一步把效果改坏了。这套流程看着笨，但比事后回滚省时间。希望帮到你。
本文还有配套的精品资源，点击获取
