简介：本资源面向具备一定深度学习基础的计算机视觉学习者与开发者，聚焦PyTorch框架下UNet与FCN两种经典图像语义分割算法的完整实现与源码解析，可用于课程设计、科研复现或工程入门。压缩包共18个文件，约227KB，以py脚本、ipynb交互式笔记为主，辅以png预测可视化图、zbak备份文件及license、md说明文档，覆盖模型定义、损失函数、训练与推理全流程。项目采用模块化设计，将流程拆解为数据预处理、模型搭建、训练优化与预测验证四环节：UNet以对称编码器-解码器配合跳跃连接实现多尺度特征融合，FCN通过全卷积化支持任意尺寸图像端到端分割，并引入混合精度训练、余弦退火与类别权重平衡等技巧。评估涵盖交并比、像素准确率等指标，代码含类型注解与单元测试，便于二次开发。目前已有81人学习下载，适合希望系统掌握分割算法细节与工程规范的读者参考。
1. 从一张广告牌说起：UNet 与 FCN 到底能切出什么
去年帮朋友处理一批户外广告牌巡检图，需求很朴素：把画面里的广告牌区域抠出来，剩下的天空、树木、行人全部归为背景。我一开始想用目标检测框一下了事，结果发现广告牌经常被灯杆、树枝遮挡，矩形框里混进大量无关像素，后续做尺寸测量和内容识别时误差大得离谱。这就是语义分割要解决的问题——不是画框，而是给每个像素分类。PyTorch 生态里，FCN 和 UNet 是两条最常被拿来落地的路线：FCN 用全卷积替换全连接，把分类网络改造成逐像素预测；UNet 在此基础上加了编码器-解码器之间的跳跃连接，让小目标和不规则边缘的定位精度明显提升。这份资源把两个算法的实现、训练脚本和推理入口都放在一起，适合已经装好 PyTorch、想跑通第一个分割项目的人，也适合需要对照源码理解跳跃连接、上采样和损失函数怎么配合的熟手。下面我按自己拆包复现的顺序，把能抄的步骤和容易翻车的地方一次讲清。
2. 环境与数据准备：把 PyTorch 和分割数据集先对齐
2.1 为什么分割项目对环境版本更敏感
图像分割和普通分类任务不同，它同时依赖卷积、上采样、转置卷积和逐像素损失，这些算子在 CUDA 版本、cuDNN 版本和 PyTorch 小版本之间偶发行为差异。我遇到过同一份 UNet 代码在 PyTorch 1.12 上正常收敛，换到 2.0 后转置卷积输出尺寸差一个像素，导致标签和预测对不齐，损失直接飙到 nan。常见做法是锁定一个经过验证的组合，比如 Python 3.9 + PyTorch 2.0.1 + CUDA 11.8，或者直接用 conda 创建独立环境，避免和系统里已有的 pytorch 基础框架互相污染。如果你在 WSL 里搭环境，注意把显卡驱动装在 Windows 侧，WSL 内只装 CUDA toolkit 和 PyTorch，否则会出现 torch.cuda.is_available() 返回 False 的玄学问题。
# 创建独立环境，避免和已有 pytorch 环境冲突 conda create -n unet_seg python=3.9 -y conda activate unet_seg # 安装 PyTorch，以 CUDA 11.8 为例，具体命令按官网当前推荐调整 pip install torch==2.0.1 torchvision==0.15.2 --index-url https://download.pytorch.org/whl/cu118 # 验证 GPU 是否可用 python -c "import torch; print(torch.__version__, torch.cuda.is_available())"上面三条命令里，第一条建环境，第二条装框架，第三条做自检。参数上唯一需要你改的是 CUDA 版本号，它必须和本机驱动支持的最高版本匹配，不是越高越好。如果输出是 False，先别急着重装，用 nvidia-smi 看驱动版本，再对照 PyTorch 官网的兼容表，多数情况是装成了 CPU 版。
2.2 数据目录怎么组织才能直接喂给 Dataset
这份源码默认读取的是 VOC 格式的分割数据集，目录结构要求比较固定。我一般会先把自己的数据整理成下面这样，再改 Dataset 里的路径，而不是反过来去改代码逻辑。
dataset/ ├── JPEGImages/ # 原图，jpg 或 png ├── SegmentationClass/ # 标签图，png，像素值就是类别 id └── ImageSets/ └── Segmentation/ ├── train.txt # 每行一个文件名，不带扩展名 └── val.txt标签图这里有个血泪经验：很多人用 labelme 或 ps 导出标签时，保存成了 RGB 三通道图，每个类别颜色不同，但像素值不是 0、1、2 这样的连续 id。UNet 训练时做交叉熵会直接拿像素值当类别索引，结果就是类别对不上、损失不下降。正确做法是确保标签图是单通道 8 位，背景为 0，目标从 1 开始递增。如果手上只有彩色标签，写个映射脚本转一下，比在训练里硬编码颜色表可靠得多。
import numpy as np from PIL import Image # 把 RGB 彩色标签转成单通道类别 id 图 color_map = { (0, 0, 0): 0, # 背景 (255, 0, 0): 1, # 类别1 (0, 255, 0): 2, # 类别2 } def rgb_to_label(mask_path, save_path): rgb = np.array(Image.open(mask_path).convert("RGB")) label = np.zeros(rgb.shape[:2], dtype=np.uint8) for color, idx in color_map.items(): match = np.all(rgb == color, axis=-1) label[match] = idx Image.fromarray(label).save(save_path) rgb_to_label("dataset/SegmentationClass/0001.png", "dataset/SegmentationClass/0001_label.png")这段脚本的核心是 color_map 字典，键是标签图里实际出现的 RGB 值，值是训练时使用的类别 id。跑之前先用取色工具确认几个像素点的颜色，别凭肉眼猜。转换完抽查几张，用 numpy.unique 看下像素值分布，确认没有遗漏的颜色被归到背景。
3. FCN 实现拆解：全卷积、上采样与跳级融合怎么落地
3.1 从分类网络到逐像素预测的关键改动
FCN 的思路是把 VGG 或 ResNet 最后的全连接层换成卷积层，让网络输出从一维向量变成二维特征图，再通过上采样恢复到原图尺寸。源码里 FCN 的主干通常用预训练 VGG16，前 13 层卷积保持不变，后面接 1x1 卷积把通道数压到类别数，最后用转置卷积放大 32 倍。这个 32 倍版本就是 FCN-32s，边缘比较粗糙；源码里还实现了 FCN-16s 和 FCN-8s，分别把 pool4 和 pool3 的特征拿过来做融合，小目标分割效果会好一些。选哪个版本取决于你的目标尺寸：广告牌、道路这类大区域，FCN-32s 够用；细胞、裂缝这类细长目标，直接上 FCN-8s，别在 32s 上浪费时间调参。
import torch import torch.nn as nn import torchvision.models as models class FCN8s(nn.Module): def __init__(self, num_classes): super().__init__() vgg = models.vgg16(pretrained=True).features # 取 VGG 不同阶段的特征，用于后续跳级融合 self.stage1 = vgg[:17] # 到 pool3 self.stage2 = vgg[17:24] # 到 pool4 self.stage3 = vgg[24:31] # 到 pool5 self.score_pool4 = nn.Conv2d(512, num_classes, 1) self.score_pool3 = nn.Conv2d(256, num_classes, 1) self.upscore2 = nn.ConvTranspose2d(num_classes, num_classes, 4, stride=2, padding=1) self.upscore8 = nn.ConvTranspose2d(num_classes, num_classes, 16, stride=8, padding=4) def forward(self, x): p3 = self.stage1(x) p4 = self.stage2(p3) p5 = self.stage3(p4) score = self.score_pool4(p4) + self.upscore2(self.score_pool3(p3)) return self.upscore8(score)这段代码里，stage1、stage2、stage3 把 VGG 切成三段，分别拿到 pool3、pool4、pool5 的特征。score_pool4 和 score_pool3 是两个 1x1 卷积，把通道数统一到类别数，方便相加。upscore2 做 2 倍上采样，upscore8 做 8 倍上采样，最终输出和输入同尺寸。参数上 num_classes 要和你标签里的类别数一致，包含背景。转置卷积的 kernel_size、stride、padding 三个值必须配套，改一个就要重算输出尺寸，否则拼接时会报维度不匹配。
3.2 训练循环里损失函数和指标怎么选
分割任务最常用的损失是交叉熵，但类别极不平衡时（比如背景占 90%），交叉熵会被背景主导，模型学会全预测背景也能拿到高准确率。源码里默认用 CrossEntropyLoss，我一般会改成带权重的版本，或者叠加 Dice Loss。指标上别只看 pixel accuracy，它在这个场景下会骗人，重点看 mIoU 和每个类别的 IoU。
import torch.nn.functional as F def train_one_epoch(model, loader, optimizer, device, weight=None): model.train() total_loss = 0 for imgs, labels in loader: imgs, labels = imgs.to(device), labels.to(device).long() optimizer.zero_grad() outputs = model(imgs) # weight 用于给少数类别更高权重，缓解不平衡 loss = F.cross_entropy(outputs, labels, weight=weight) loss.backward() optimizer.step() total_loss += loss.item() return total_loss / len(loader)weight 是一个长度为 num_classes 的张量，值越大表示该类越重要。计算方式常见的是取类别频率的倒数，再归一化。如果你不想手算，先用默认交叉熵跑几轮，看验证集里小类别的 IoU 是不是接近 0，是的话再加权重。学习率方面，FCN 微调预训练主干时用 1e-4 比较稳，从零训练可以到 1e-3，但 batch size 要相应调整。
4. UNet 实现拆解：编码器-解码器与跳跃连接的真实作用
4.1 下采样、上采样和拼接的尺寸对齐
UNet 的结构像一个 U 形：左边编码器不断卷积加池化，特征图变小、通道变多；右边解码器不断上采样加卷积，特征图变大、通道变少；中间用跳跃连接把编码器同层的特征直接拼到解码器对应层。这个拼接是 UNet 比 FCN 在医学图像分割、小目标分割上表现更好的核心原因，因为它保留了下采样过程中丢失的空间细节。但拼接对尺寸要求很严格，编码器第 n 层输出是 64x64，解码器对应层上采样后也必须是 64x64，差一个像素就报错。我一般会在拼接前打印两边 shape，确认一致再往下走。
class UNet(nn.Module): def __init__(self, in_channels=3, num_classes=2): super().__init__() def block(in_c, out_c): return nn.Sequential( nn.Conv2d(in_c, out_c, 3, padding=1), nn.BatchNorm2d(out_c), nn.ReLU(inplace=True), nn.Conv2d(out_c, out_c, 3, padding=1), nn.BatchNorm2d(out_c), nn.ReLU(inplace=True), ) self.enc1 = block(in_channels, 64) self.enc2 = block(64, 128) self.pool = nn.MaxPool2d(2) self.bottleneck = block(128, 256) self.up2 = nn.ConvTranspose2d(256, 128, 2, stride=2) self.dec2 = block(256, 128) # 拼接后通道翻倍，所以输入是 256 self.up1 = nn.ConvTranspose2d(128, 64, 2, stride=2) self.dec1 = block(128, 64) self.out = nn.Conv2d(64, num_classes, 1) def forward(self, x): e1 = self.enc1(x) e2 = self.enc2(self.pool(e1)) b = self.bottleneck(self.pool(e2)) d2 = self.up2(b) d2 = self.dec2(torch.cat([e2, d2], dim=1)) # 跳跃连接 d1 = self.up1(d2) d1 = self.dec1(torch.cat([e1, d1], dim=1)) return self.out(d1)代码里 enc1、enc2 是编码器两个阶段，bottleneck 是底部，up2、up1 是转置卷积上采样，dec2、dec1 是解码器卷积块。关键点在 dec2 的输入通道是 256，因为 torch.cat 把 e2 的 128 通道和上采样后的 128 通道拼在一起了。如果你把 num_classes 改成 3 或更多，只需要改 out 层，其他不用动。BatchNorm 在小 batch size 下表现不稳定，如果显存只够跑 batch size 2，建议换成 GroupNorm 或 InstanceNorm。
4.2 训练自己的数据集要改哪几个地方
拿到源码后直接跑 demo 数据通常没问题，但换成自己的数据，至少要改三处：Dataset 里的路径和类别数、模型初始化时的 num_classes、以及可视化时的颜色表。我见过有人只改了路径没改类别数，训练不报错但预测全是背景，排查半天才发现输出通道还是 2。另外，UNet 对输入尺寸没有硬性要求，但为了下采样和上采样能整除，建议把图片 resize 到 16 的倍数，比如 256x256 或 512x512。如果原图长宽比很重要，用 padding 而不是直接拉伸。
from torch.utils.data import Dataset, DataLoader from PIL import Image import torchvision.transforms as T class SegDataset(Dataset): def __init__(self, img_dir, mask_dir, file_list, size=256): self.img_dir = img_dir self.mask_dir = mask_dir self.names = [l.strip() for l in open(file_list)] self.size = size self.img_tf = T.Compose([T.Resize((size, size)), T.ToTensor()]) self.mask_tf = T.Compose([T.Resize((size, size), interpolation=T.InterpolationMode.NEAREST)]) def __len__(self): return len(self.names) def __getitem__(self, idx): name = self.names[idx] img = Image.open(f"{self.img_dir}/{name}.jpg").convert("RGB") mask = Image.open(f"{self.mask_dir}/{name}.png") return self.img_tf(img), torch.from_numpy(np.array(self.mask_tf(mask))).long() loader = DataLoader(SegDataset("dataset/JPEGImages", "dataset/SegmentationClass", "dataset/ImageSets/Segmentation/train.txt"), batch_size=4, shuffle=True, num_workers=2)标签图的 resize 必须用最近邻插值，用双线性会把类别 id 插成小数，再取整就乱了。num_workers 在 Windows 上有时会卡死，设成 0 先跑通，再逐步往上加。batch_size 根据显存调，8G 显存跑 256x256 的 UNet，batch size 4 到 8 比较合适。
5. 避坑与排查：分割训练里最容易翻车的五件事
5.1 损失不下降，先查标签而不是模型
现象：训练几个 epoch，loss 在 0.69 附近震荡，准确率不动。原因：标签图里只有 0 和 255 两个值，255 被当成类别 255，而模型输出只有 2 类，交叉熵计算时目标越界，实际梯度是乱的。解决：用 numpy.unique 检查标签像素值，确保是 0 到 num_classes-1 的连续整数，255 要映射成 1。
5.2 验证集 IoU 很高但预测图全黑
现象：mIoU 显示 0.9，但把预测结果叠到原图上看，目标区域根本没被标出来。原因：背景占比太高，模型全预测背景，背景 IoU 拉高了平均值，小类别 IoU 是 0。解决：分开打印每个类别的 IoU，别只看平均值；同时把验证集里目标占比高的样本单独抽出来看。
5.3 转置卷积输出尺寸和输入对不上
现象：UNet 拼接时 RuntimeError: Sizes of tensors must match。原因：输入图片尺寸不是 16 的倍数，经过多次下采样后奇数尺寸被取整，上采样回来差一个像素。解决：在 Dataset 里统一 resize 到 16 的倍数，或者在拼接前用 F.interpolate 把上采样结果对齐到编码器特征的尺寸。
5.4 显存够但训练速度极慢
现象：GPU 利用率只有 20%，一个 epoch 要跑十几分钟。原因：num_workers 设成 0，数据加载在主进程串行执行，GPU 一直在等数据。解决：把 num_workers 调到 CPU 核心数的一半左右，同时开 pin_memory=True，让数据搬运和计算重叠。
5.5 推理结果比训练时差很多
现象：训练时指标正常，用单张图推理时边缘破碎、类别错乱。原因：推理时忘了加 model.eval()，BatchNorm 还在用当前 batch 的统计量；或者输入没有做和训练一致的归一化。解决：推理前固定写 model.eval() 和 torch.no_grad()，并把训练时的均值和标准差抄过来做 Normalize。
6. 进阶技巧：把 UNet 推理结果导出 ONNX 并验证一致性
训练完模型只是第一步，真正部署时经常要把 PyTorch 模型转成 ONNX，再交给推理引擎或板端。这一步最容易出的问题是导出成功但结果对不上，所以我会强制做一次数值比对。下面这段脚本把 UNet 导出成 ONNX，然后用 onnxruntime 跑同一张输入，和 PyTorch 输出比最大绝对误差。
import torch import numpy as np import onnxruntime as ort model = UNet(num_classes=2) model.load_state_dict(torch.load("best_unet.pth", map_location="cpu")) model.eval() dummy = torch.randn(1, 3, 256, 256) torch.onnx.export(model, dummy, "unet.onnx", input_names=["input"], output_names=["output"], opset_version=11, dynamic_axes={"input": {0: "batch"}}) with torch.no_grad(): torch_out = model(dummy).numpy() sess = ort.InferenceSession("unet.onnx") onnx_out = sess.run(None, {"input": dummy.numpy()})[0] print("max abs diff:", np.abs(torch_out - onnx_out).max())导出时 opset_version 建议用 11 或更高，低版本对某些上采样算子支持不好。dynamic_axes 把 batch 维设成动态，部署时可以传不同 batch size。比对结果里 max abs diff 在 1e-4 以内算正常，超过 1e-2 就要查是不是有算子被替换成了近似实现。我一般还会把 ONNX 模型的输出 argmax 成类别图，和 PyTorch 的 argmax 结果逐像素比，确认没有大面积翻转。这套流程走完，模型才算真正能交出去。从那以后我每次导出 ONNX 都强制跑一遍数值比对，不再凭“导出没报错”就认为没问题。希望帮到你。
本文还有配套的精品资源，点击获取
