简介一份基于PyTorch实现SegNet图像分割任务的Python源码项目出自作者经导师指导并获98分的高分期末大作业主要面向计算机相关专业正在完成课程设计、期末大作业的学生以及需要图像分割实战练习的开发者。压缩包共119个文件约27.19MB包含14个Python脚本、77个图像样本、1个pth模型权重文件另有sh启动脚本、Dockerfile、README说明、日志与配置文件完整覆盖从环境配置到模型训练与推理的流程。目前已有176人学习浏览。这套源码可直接运行结合代码与文档可理解编码器-解码器、池化索引复用等核心机制也可作为课程设计报告的项目范例便于在此基础上进行二次开发与实验扩展。1. 图像分割大作业为什么我先把宝押在 SegNet 上第一次看到这个题目大多数人第一反应是 U-Net 或者 FCN毕竟讲分割的课总是把这两个挂在嘴边。但如果你认真翻过这几年的课程大作业选题会发现一个很微妙的规律好写的模型成果卷不动不好写的模型反而容易出东西。SegNet 恰好是那个「答辩时能让老师停下来问你原理」的选项。它在语义分割上做的事很简单——把 VGG16 的分类骨架拆成编码器和解码器用 max-pooling 索引代替反卷积做上采样但正是这个「索引复用」的设计让它在参数量、显存占用和边界保留上有一套完全不同的取舍。多数学生第一次跑通它靠的是一份能直接落地的 PyTorch 源码而不是从零堆网络。这份基于 PyTorch 实现 SegNet 的图像分割任务源码适合两类人一是课程大作业需要快速产出可演示结果的同学二是想弄清楚「池化索引上采样到底比反卷积好在哪」的入门研究者。接下来我会按工程落地的顺序把结构、数据、训练、避坑和交付一次讲透。2. SegNet 在学什么结构拆解与选型理由2.1 编码器部分VGG16 骨架和它到底保留了哪些东西SegNet 的编码器基本是把 VGG16 的前 13 层搬过来分成五个编码块。前两块各两层卷积后三块各三层卷积每块结束接一个 2×2 的 max-pooling。这里最关键的一步不是卷积本身而是 pooling 时把「最大值所在的位置」记下来这个位置索引在后面解码时会原封不动用回去。你可能要问池化索引不就省了反卷积吗效果能差多少实际差别体现在训练稳定性和边界质量上。反卷积的卷积核是学出来的学不好会出现棋盘格伪影池化索引不学任何参数上采样后边界位置天然对齐。我在落地这份源码时最在意的其实是显存。SegNet 的编码器每一层都不算大但 PyTorch 里如果你直接复用 torchvision 的 VGG16会带上分类头那三层全连接白白占几十 MB 不说还会拖慢加载速度。常见做法是只用 features 部分并且去掉最后一个 maxpool 之前的输出对齐逻辑让编码器输出分辨率变成原图的 1/32。有一个细节容易被忽略VGG16 的 features 里第 4 和第 5 个卷积块原版用的是 padding1 的 3×3 卷积如果你为了「更快」把 padding 改成 0后面解码器的索引对齐会乱feature map 尺寸差 2 个像素小数据集上不报错但 loss 曲线会一直震荡。import torch import torch.nn as nn def make_encoder(vgg16_features): # 取 VGG16 的 features并按 SegNet 的 5 个编码块切分 blocks [] stages [2, 2, 3, 3, 3] # 每块卷积层数 idx 0 for stage in stages: block [] for _ in range(stage): block.append(vgg16_features[idx]) idx 1 # 每个编码块后面跟一个带索引记录的 MaxPool blocks.append(nn.Sequential(*block)) idx 1 # 跳过原 VGG16 的 MaxPool后面自定义 return nn.ModuleList(blocks)这段代码把预训练 VGG16 的 features 按 SegNet 论文的划分方式切成五个子模块原版 pooling 全部跳过这样每个编码块后面就可以接自定义的带索引 MaxPool2d。注意这里是用 ModuleList 包起来保证模型在 .cuda() 和 .train() 时能正确切换状态。参数层面的好处是只有卷积层参与计算分类头直接丢弃加载预训练权重时用 strictFalse只拿能对上的 key省去一堆维度报错。2.2 解码器部分池化索引上采样与卷积修复的组合解码器是做「反着走」的。每个解码块先做一次 max-unpooling用编码器存的索引把当前 feature map 里的值放回到上采样前的位置其余位置补 0。这一步会让分辨率翻倍但特征图会变稀所以紧接着要接一层卷积把稀疏的特征稠密化。论文里解码块和编码块是镜像对称的卷积核数量从 512 一路降到 64最后一层用 1×1 卷积输出类别数。一个容易翻车的实现细节是 max-unpooling 的索引格式。PyTorch 的 MaxPool2d 返回 indices但如果你把 indices 直接传给 MaxUnpool2d要保证两次调用时的 kernel_size、stride 和 padding 完全一致否则维度对不上。另外如果你把模型放到 GPU 上跑indices 是 int64 类型的 Tensor和 float 权重混在一个 list 里做前向传播没问题但如果你想自己拼接 torch.cat 就会报类型错误。class DecoderBlock(nn.Module): def __init__(self, in_ch, out_ch, num_convs2): super().__init__() self.unpool nn.MaxUnpool2d(kernel_size2, stride2) convs [] for _ in range(num_convs): convs.append(nn.Conv2d(in_ch, in_ch, 3, padding1)) convs.append(nn.BatchNorm2d(in_ch)) convs.append(nn.ReLU(inplaceTrue)) convs.append(nn.Conv2d(in_ch, out_ch, 3, padding1)) self.convs nn.Sequential(*convs) def forward(self, x, indices): x self.unpool(x, indices) x self.convs(x) return x前向时要手动把编码器存下来的 indices 逐个传进来这是 SegNet 实现里最容易写乱的地方。我的习惯是在编码器每个 MaxPool 后面直接把 indices 存到一个列表里而不是返回一个字典省去键名错配的时间。同时注意解码器最后一层不要加 ReLU因为后面要接 CrossEntropyLoss它会自己算 log_softmax输出层再加激活会双重缩放。2.3 选型对比SegNet 和 FCN、U-Net 在课程作业里的真实差异课程作业跟论文复现不一样判分重点往往落在「结构清晰、能跑通、有对比实验」。SegNet 和 FCN 的长短板非常互补FCN 的反卷积模块参数多但表达能力强适合小目标SegNet 的解码器参数少训练更快边界锐利度更好但对非常细小的结构容易漏。U-Net 用 skip connection 把编码器特征直传解码器小数据集上效果最好缺点是显存占用高而且「跳连接」这个点在答辩时几乎人人都会问反而不好讲出新意。对比项FCNU-NetSegNet上采样方式反卷积反卷积跳连池化索引解码器参数量高高低训练速度快慢很快边界质量一般很好好答辩发挥空间小小大表格只是参考真正落地时建议把这三种模型都跑一遍哪怕只是每个模型训 100 个 iter 看 loss 下降速度也能给你写「选型理由」这段报告攒素材。SegNet 不是每个任务的最优解但作为大作业它是一个「参数少、易解释、表现稳定」的稳妥解。3. 拿到源码先拆结构文件布局与数据准备3.1 工程文件清单与每个文件负责的事这类打包源码通常不会只有一个 .py 文件。最典型的工程结构是五个文件模型定义、数据集读取、工具函数、训练入口、预测可视化。我拿到压缩包的第一件事不是跑 train.py而是先看模型定义文件里有没有预训练权重加载逻辑。很多同学踩过的坑是模型定义和数据集读取都写得很完整但训练入口里没有设置随机种子导致同一个代码每次跑出来的 mIoU 差两三个点这在课程作业的「复现实验」环节非常吃亏。文件常见命名主要职责模型结构segnet.py / model.pyEncoder、Decoder、完整 SegNet 类数据读取dataset.py / loader.py读取图像与标注、数据增强、返回 batch工具函数utils.py混淆矩阵、mIoU 计算、学习率调整训练脚本train.py主流程、参数解析、日志输出推理可视化predict.py / demo.py加载权重、对单张图预测、着色保存检查完文件清单下一步是确认标注读取方式。VOC 格式的标注是调色板 PNG不是普通的 RGB 三通道图如果你用 Image.open(img_path).convert(RGB) 去读得到的三通道值并不等于类别 ID后面计算 loss 会全错。正确做法是用 convert(P) 读成调色板模式把像素值直接当作类别索引。这一步错了训练不会报错但 mIoU 会长期徘徊在极低水平。3.2 把数据整理成 SegNet 需要的目录结构SegNet 的数据加载常见做法是维护一个 image / label 路径一一对应的列表而不是强制要求某个固定目录。所以第一步是把数据集拆成两张清单train.txt 和 val.txt每行是一对路径用空格或逗号分隔。写一个小脚本自动扫描目录既避免手写路径出错也方便后续替换成自己的数据。import os def build_split(images_dir, labels_dir, val_ratio0.15): imgs sorted(os.listdir(images_dir)) lines [] for name in imgs: base os.path.splitext(name)[0] label_path os.path.join(labels_dir, base .png) if os.path.exists(label_path): lines.append(f{os.path.join(images_dir, name)} {label_path}\n) val_n max(1, int(len(lines) * val_ratio)) with open(train.txt, w, encodingutf-8) as f: f.writelines(lines[val_n:]) with open(val.txt, w, encodingutf-8) as f: f.writelines(lines[:val_n])这个脚本做的事就是扫描两个目录自动把名字相同的图像和标注配对然后按比例切出验证集。切分时直接用列表切片把后面的样本给训练集前面的给验证集。注意这里没有做随机打乱如果想打乱需要在切片前加 random.shuffle。真实使用中标签后缀可能是 .png 也可能是 .jpg要根据数据集实际情况改这行代码否则会漏掉一批样本。3.3 先做一个小数据集冒烟测试不要直接上全量全量 VOC 训练一次要几个小时如果你代码里有 bug等到 loss 曲线异常才反应过来时间成本太高。我一般会先从训练集里随机抽出 20 张图片单独建一个 smoke 文件夹把 crop size 调小到 224batch size 设成 2只训 50 个 iter确认 loss 在下降、显存没爆、验证脚本能跑通再切回全量数据训练。# 在 train.py 里加一个 --smoke 参数快速截取子集用于调试 if args.smoke: dataset torch.utils.data.Subset(dataset, range(20)) args.max_iter min(args.max_iter, 50)用 Subset 包裹原数据集是最省事的截断方式不需要改动 Dataset 类的代码也不会影响后续全量训练的数据加载逻辑。训练迭代数用 min 兜底避免忘改参数导致直接跑完整个训练流程。这个习惯能帮你区分「代码 bug」和「模型效果差」两类问题后者需要看曲线趋势前者只需要看 loss 数值是否在下降到合理范围。4. 训练与验证从命令行参数到 mIoU 指标4.1 训练主入口与必调参数训练脚本的正常打开方式不是直接双击跑而是用命令行参数控制超参数。源码里通常会有一堆 argparse 参数但大作业场景下真正需要动的不超过六个batch size、base learning rate、max iter、crop size、num classes、backbone 预训练路径。我在给 A 同学调代码时最常改的就是 crop size 和 batch size 的搭配因为这直接决定显存够不够用。python train.py \ --backbone vgg16_bn \ --dataset voc \ --train-list train.txt \ --val-list val.txt \ --batch-size 8 \ --base-lr 0.01 \ --max-iter 20000 \ --crop-size 384 \ --num-classes 21 \ --gpu 0参数含义看名字基本能猜到但有几个容易误解的地方。backbone 选 vgg16_bn 会在卷积后带 BatchNorm训练更稳定代价是前向速度略慢base-lr 0.01 是配合 poly 学习率衰减策略的初始值如果源码用的是 step 衰减0.01 就会偏大loss 容易炸。crop size 384 是经验值显存 8 GB 的卡上配 batch size 8刚刚好卡在边界再往上加就会 OOM。训练时的日志输出至少要包含三类信息当前 iter、loss 数值、当前学习率。如果源码没有打印学习率强烈建议自己加一行因为学习率变化直接影响 loss 曲线的形态答辩时被问到「为什么 loss 在 15000 iter 处有拐点」如果没有学习率的佐证很难解释清楚。4.2 损失函数与类别不均衡的两个处理手法图像分割默认用 CrossEntropyLoss 逐像素分类。但 VOC 这类数据集中背景像素占比经常超过 70%模型只需要把所有像素都预测成背景mIoU 就能到 60% 上下这是典型的类别不均衡问题。常见的处理办法有两个一是给 loss 加 class weight让低频类别获得更大梯度二是给 loss 设置 ignore_index把标注里边缘的「硬像素」或 255 类别的像素直接跳过。criterion nn.CrossEntropyLoss( weightclass_weight, ignore_index255, )class_weight 的计算方式一般取每个类别像素频率的倒数再开根号平滑避免权重过大导致低频类别过拟合。我不建议直接使用整个数据集的统计频率而是用每个 batch 内统计的近似值训练开始时波动会大一些但 1000 个 iter 后会稳定下来。ignore_index255 看起来是小改动实际影响很大因为标注里的边界、难样本和部分数据集填充值都落在 255 上不忽略的话loss 会被这些既难学又不参与评估的像素带着走。参数设置完毕后训练曲线会有两种典型形态。第一种是 loss 快速下降并趋平mIoU 稳步上升这是正常形态第二种是 loss 在 100 iter 内降到很低但 mIoU 一直不动这时十有八九是模型把全部像素预测成了背景类。确认方法很简单取一个 batch 的预测结果打印 argmax 后各类别的像素占比如果 99% 都落在一个类上就说明学习率太大或类别权重设置错误。4.3 预测与可视化把像素索引变回彩色图训练完的权重放在 .pth 文件里预测脚本的核心逻辑是加载权重、前向计算、取 softmax 后最大概率的类别索引、用调色板映射成彩色图。VOC 的调色板是固定的 256 色但每个类别只用了其中一个索引所以需要先建一个颜色列表再按类别索引取色。import numpy as np import torchvision.transforms.functional as F def decode_segmap(pred, num_classes21): # pred: tensor shape (1, H, W)值为类别索引 label_colors np.zeros((num_classes, 3), dtypenp.uint8) # 这里按常规 VOC 配色初始化也可改为自定义颜色 label_colors[0] [0, 0, 0] label_colors[1] [128, 0, 0] # 省略中间类别颜色填充 r label_colors[pred, 0] g label_colors[pred, 1] b label_colors[pred, 2] return np.stack([r, g, b], axis2)这个函数把 H×W 的类别索引数组映射成 H×W×3 的彩色图。注意 pred 必须先在 CPU 上转成 numpy 数组再索引直接在 GPU tensor 上做这种操作会拖慢速度且容易出错。可视化脚本还有一个容易被忽略的功能点原图和预测图要拼在一起保存最好再叠加上半透明的预测结果这样生成的图片在报告里一眼就能看出差别比单独放两张图更有说服力。预测阶段的数据预处理要和训练时完全一致尤其是归一化的 mean 和 std。如果训练时用的是 ImageNet 统计值 [0.485, 0.456, 0.406]预测时也一定要用同一个归一化否则预测图会整体偏暗或发灰肉眼看着像模型训坏了实际只是颜色通道被缩放错误这类「假翻车」在答辩前最容易让人白紧张一场。4.4 验证阶段mIoU 和 PA 怎么算才不会被答辩老师挑刺mIoU 是语义分割最常用的指标计算方式是先统计混淆矩阵然后对每个类别算 IoU最后取平均。PA 是像素准确率即正确分类的像素数除以总像素数模型把所有像素都预测成背景时 PA 可能高达 80%但 mIoU 会很低。答辩时两个指标要一起给并且要解释清楚为什么 PA 高不代表效果好。def compute_miou(hist): with np.errstate(divideignore, invalidignore): iou np.diag(hist) / (hist.sum(axis1) hist.sum(axis0) - np.diag(hist)) iou iou[np.isfinite(iou)] return np.mean(iou)hist 是 num_classes × num_classes 的混淆矩阵维护方式是在验证循环里对每个像素的预测值和标签累加。注意分母里 hist.sum(axis1) 加 hist.sum(axis0) 再减对角线这是 IoU 的标准定义交集除以并集。用 np.errstate 忽略除零警告是因为某些类别可能在验证集里一次都没出现完全不忽略的话控制台刷一堆 RuntimeWarning影响日志阅读。验证时要严格控制训练模式和验证模式的切换。模型在 .train() 和 .eval() 模式下BatchNorm 和 Dropout 的行为完全不同尤其是 BatchNormeval 模式使用的是统计均值而不是当前 batch 的均值。很多同学训练完 mIoU 看着正常一跑验证就掉点甚至 loss 出现负数多半就是忘了切模型状态。5. SegNet 训练最常见的 5 个坑现象、原因与解决5.1 显存不够crop size 和 batch size 的组合要压到卡能承受的范围现象训练脚本跑到第一个迭代就报 CUDA out of memory或者跑了几百个 iter 后在某个 batch 突然崩溃。原因有两种一是 crop size 设置太大比如直接 512二是 batch size 和 crop size 的乘积超出了显存容量。解决路径是先把 batch size 降到 2再逐步降 crop size直到能稳定跑完 50 个 iter同时注意 PyTorch 的显存分配机制tensor 的临时缓冲区会在反向传播时翻倍。有个提升利用率的技巧是把 backbone 冻结前几层让这部分不参与反向传播节省的显存可以用来支撑更大的 crop size。5.2 mIoU 一直在 0.1 徘徊先检查是不是全预测成了背景现象loss 在下降但验证集 mIoU 长期低于 0.2可视化结果全是背景色。原因几乎都是类别不均衡背景像素占比过高模型找到了「全部预测为背景」这个局部最优解。解决时先往 CrossEntropyLoss 里加 class_weight用小数据集重新训练确认输出不再单调后再逐步拉开权重。这个坑最玄学的地方在于不同随机种子下复现结果不一样有的种子下模型能自己跳出局部最优有的就陷在里面出不来所以一定要固定随机种子否则你没法判断是代码问题还是运气问题。5.3 BatchNorm 在 batch size 很小时翻车表现是 loss 曲线锯齿状现象loss 曲线整体下降但每个 step 都剧烈抖动验证精度也随之波动。原因是 batch size 太小BatchNorm 算出来的均值和方差噪声极大模型每步都在适应不同的数据分布。解决方式有两个一是把 batch size 提到至少 4二是如果显存实在不够就改用 GroupNorm 替代 BatchNorm。替换方法是把模型里所有 BatchNorm2d 换成 GroupNorm(num_channels32)不改变模型整体结构效果在大作业这种数据规模下往往比原版更稳定。5.4 验证预测图画出来整体发黑归一化和可视化脱节现象模型明明有不错的 mIoU 指标但保存出来的预测图和原图颜色不在一个亮度区间看起来像全黑或褪色。原因是训练时的图像做了归一化而可视化时没有反归一化就直接保存预览图。解决方法是可视化前先乘以 std 再加 mean把像素值恢复到 0-255 范围再转 uint8 保存。这个坑跟模型效果无关纯属图像处理流程的小细节遇到时先用 matplotlib 在调试窗口显示一下确认预览图正常再保存。5.5 Windows 下 dataloader 多进程崩溃num_workers 直接设 0现象代码在 Linux 上跑得好好的换到 Windows 上运行几个 epoch 后控制台报 worker 相关错误或者程序直接卡死没有输出。原因是 Windows 的进程启动机制和 Linux 的 fork 不同PyTorch DataLoader 多进程在 Windows 上更容易出现主模块重复导入的问题。解决方式是 num_workers 设成 0 或者把数据加载逻辑放进一个单独的 Python 脚本里。这个坑在课程作业阶段几乎必踩尤其是同学之间用不同操作系统互跑代码时经常被误判为「代码本身有 bug」实际上只是多进程兼容性问题。6. 交付与进阶让「高分」落到报告和复现里源码跑通只是大作业的下半场真正拉开分数差距的是「可复现演示」和「实验记录」。我交作业前一定会做三件事固定随机种子并打印到日志、用同一份权重重新跑一次完整验证、把所有训练日志和可视化结果放进同一个文件夹。操作起来很简单在 train.py 开头加三行代码random.seed(0) np.random.seed(0) torch.manual_seed(0)在数据加载器里也要对应设置 generator 的随机种子否则数据增强的随机性会让每次验证结果有微小差异。如果你想让结果更稳把跑完的训练日志和损失曲线截图一起归档这样答辩被问到「这个效果是你哪一次实验跑出来的」你可以直接翻出对应日志而不是支支吾吾说「好像是某一次」。进阶用法是在现有的 SegNet 上做一个小改动并写进报告比如把解码器的普通卷积替换成带空洞卷积的版本在保持参数量基本不变的情况下增大感受野专门处理大目标边缘粗糙的问题。这个改动只要两行代码但能让你在「创新点」这一栏有话可说。另一个性价比很高的改动是把验证集的每类 IoU 单独打印出来找出模型最弱的三个类别然后在报告里针对其中一个做一个小实验比如给对应类别的 loss 提高权重。这种「发现问题—分析原因—实验验证」的闭环比堆砌十个图表更能打动评分的人。最后说一个我自己的习惯。每次交源码大作业之前我会把整个流程从头跑一遍从最原始的压缩包解压到训练结束再验证一次全程记录命令行窗口的截图。这个流程常常会发现「环境配置不同导致跑不通」这类问题这些问题在功能上早就验证过了但换个目录换个环境就会暴露。提前走一遍这个流程能省去交作业前一天晚上临时找老师补救的尴尬。希望这个方式能帮到你至少让你在交代码前心里有个底。本文还有配套的精品资源点击获取