简介基于 PyTorch 实现 SegNet 图像分割任务的完整工程源码面向计算机相关专业正在完成课程设计或期末大作业的学生也适合希望借助完整项目提升实战能力的学习者。项目经导师指导并获 98 分评价代码覆盖数据集加载、网络定义、训练与推理全流程可直接运行也适合作为二次开发的底稿。压缩包共 119 个文件以 PNG 图像、Python 源码、Pyc 编译文件、Shell 脚本、模型权重、Markdown 文档、Dockerfile 及配置文件等类型为主整体体积约 27.19MB。其中图像文件用于数据集样本与分割结果可视化Python 脚本实现 SegNet 核心逻辑pth 权重文件支持加载后直接测试Dockerfile 与环境配置则帮助快速搭建一致的运行环境便于训练过程复现。目录内还包含运行日志和说明文档可协助排查训练中的常见问题。项目结构按数据、模型、工具等模块划分方便按需检索与修改。目前已有 176 人学习对想参考高分项目结构或快速搭建图像分割代码框架的读者具有直接参考价值。1. SegNet不是新模型但用来做图像分割大作业依然很能打如果你的期末大作业要求是“基于深度学习实现图像分割”而你又不想在U-Net、DeepLabV3 这些被写烂的模型里跟人撞车SegNet 是一个被低估的好选择。这是一篇基于 PyTorch 实现 SegNet 的完整落地笔记从原理拆解到源码组织、从数据准备到避坑排查全部按可复现的标准来写。SegNet 的核心特点是用“池化索引”来做上采样这个机制让它在分割边界保持上比 FCN 更细腻而且结构对称、容易讲清楚非常适合答辩时演示。我见过不少学生拿着这份源码做城市道路分割、车道线检测、医学细胞分割最后拿到的评价普遍不错。这套方案适合有一点点 PyTorch 基础、想用一份完整代码完成大作业的人你不需要自己从零设计网络但需要看懂每一行在干什么。2. SegNet 原理拆解pooling indices 为什么是它的灵魂2.1 编码器-解码器对称结构从 VGG16 借来的骨干SegNet 的网络结构非常直观左边是一路编码器右边是一路解码器中间没有全连接层整体是“卷积-池化-上采样-卷积”的对称设计。编码器直接借用 VGG16 的前 13 层卷积这在 PyTorch 里一句models.vgg16_bn(pretrainedTrue).features就能拿过来省掉了从头训练特征提取器的成本。解码器做的事情是对编码器的每一层做镜像恢复每次上采样后接一个卷积层把特征通道数逐步降下来最后用 softmax 逐像素分类。SegNet 论文里给的深度是 5 组 Encoder 和 5 组 Decoder每组里的卷积层数分别是 2、2、3、3、3这正好对应 VGG16 的卷积层分布。从落地角度说Encoder 的权重用什么初始化很重要。我一般直接加载 ImageNet 上预训练好的 VGG16 权重而不是随机初始化。原因很实际分割任务普遍数据量不够大预训练权重能让模型在训练初期就有一个像样的特征表示损失下降的速度肉眼可见地快。如果你做大作业手里只有几百张图这个差别会非常明显。有个常见的误用直接把 VGG16 的所有层都拉进 Encoder包括最后的 AdaptiveAvgPool 和分类层。这是错的SegNet 的 Encoder 只需要features部分也就是卷积和池化层全连接层和池化层的索引必须自己管理。注意VGG16 的features里包含 MaxPool2d 层但 PyTorch 的features默认不返回池化索引所以不能直接拿来用。需要自己重新包一层把池化操作单独拆出来记录 indices这个在第三节代码里会详细写。2.2 MaxPool 与 MaxUnpool一对必须成对使用的算子SegNet 最核心的机制是 MaxPool 时记录下最大值所在的位置索引然后在 Decoder 里用 MaxUnpool 按照这些索引把特征“放回”原来的位置。这样做的效果是上采样时不是简单地插值而是把编码器池化时丢掉的位置信息恢复回来边界细节保留得更好。在 PyTorch 里这个机制对应两个算子torch.nn.MaxPool2d和torch.nn.MaxUnpool2d。注意MaxPool2d必须设置return_indicesTrue才会输出索引这个索引是一个 LongTensor形状和池化后的特征图一样记录的是每个池化窗口内最大值在原特征图中的位置。Decoder 里对应的MaxUnpool2d接收两个输入一个是编码器对称层传下来的索引另一个是上一层解码器输出的特征图。调用方式如下unpooled self.unpool(pooled_feat, indices)这里indices必须是编码器对应池化层保存下来的那份不能自己随便造也不能拿别层的索引替代。因为索引记录的是绝对位置而特征图的尺寸是逐层变化的用错索引轻则尺寸对不上重则直接报错。从梯度流动的角度理解一下MaxPool 在反向传播时只把梯度传给池化窗口里最大值对应的那个位置其他位置梯度为 0。MaxUnpool 是反过来只把梯度放到正向时记录的位置上。所以这两个算子天然成对中间不能插入会改变空间尺寸的操作否则就出问题。2.3 SegNet / FCN / U-Net 的取舍答辩被追问时怎么答SegNet 常被拿来和 FCN、U-Net 对比这是答辩时的高频问题。FCN 的上采样用的是转置卷积或者双线性插值参数是学出来的但上采样过程不感知编码器细节所以边界会偏模糊。U-Net 是通过 skip connection 把编码器每一层的特征拼到解码器对应层信息充足但对显存的要求更高而且结构上比 SegNet 复杂。SegNet 的定位介于两者之间没有 skip connection但通过池化索引保留了编码器池化时的位置信息。好处是解码器只需要处理编码器输出的特征不需要拼接显存占用比 U-Net 小坏处是池化索引只保留位置不保留具体值细节信息其实还是有所丢失的。从大作业的角度说选 SegNet 有三个实际理由第一结构对称清晰画网络结构图很好看答辩讲起来逻辑顺第二训练开销比 U-Net 小一张普通的 GTX 1660 就能跑起来第三和 FCN 相比有明确的创新点可以讲——就是那个 pooling indices 机制老师一问“你和 FCN 有什么区别”你直接把这个机制讲清楚就够了。如果老师追问“为什么不用 skip connection”你可以答skip connection 在医学影像和遥感图上效果好但对自然图像数据集比如 CamVid 和 VOCSegNet 的池化索引机制在保持边界锐利度和训练效率上更均衡。记住这个说法比当场支支吾吾强很多。3. 基于 PyTorch 实现 SegNet构建模型与损失函数3.1 搭建 Encoder加载 VGG16 预训练权重这里给一份可以直接跑通的 SegNet 模型代码。整个模型的实现思路是先定义 Encoder 的每一层单独处理卷积块和最大池化然后在 Decoder 里做对称展开。import torch import torch.nn as nn from torchvision import models class SegNetEncoder(nn.Module): def __init__(self, in_channels3): super().__init__() vgg16 models.vgg16_bn(pretrainedTrue) features vgg16.features # 按 VGG16 的卷积层分组每组包一个 MaxPool2d 并记录索引 self.block1 nn.Sequential( features[0], features[1], features[2], # Conv2d BN ReLU features[3], features[4], features[5], ) self.pool1 nn.MaxPool2d(kernel_size2, stride2, return_indicesTrue) self.block2 nn.Sequential( features[7], features[8], features[9], features[10], features[11], features[12], ) self.pool2 nn.MaxPool2d(kernel_size2, stride2, return_indicesTrue) self.block3 nn.Sequential( features[14], features[15], features[16], features[17], features[18], features[19], features[20], features[21], features[22], ) self.pool3 nn.MaxPool2d(kernel_size2, stride2, return_indicesTrue) self.block4 nn.Sequential( features[24], features[25], features[26], features[27], features[28], features[29], features[30], features[31], features[32], ) self.pool4 nn.MaxPool2d(kernel_size2, stride2, return_indicesTrue) self.block5 nn.Sequential( features[34], features[35], features[36], features[37], features[38], features[39], features[40], features[41], features[42], ) self.pool5 nn.MaxPool2d(kernel_size2, stride2, return_indicesTrue) def forward(self, x): indices [] x self.block1(x) x, idx1 self.pool1(x); indices.append(idx1) x self.block2(x) x, idx2 self.pool2(x); indices.append(idx2) x self.block3(x) x, idx3 self.pool3(x); indices.append(idx3) x self.block4(x) x, idx4 self.pool4(x); indices.append(idx4) x self.block5(x) x, idx5 self.pool5(x); indices.append(idx5) return x, indices这段代码有个细节要说明features里每隔一个元素就有一个MaxPool2d这些带return_indicesTrue的池化层是重新定义的不能直接复用features里的原始池化层。而且pretrainedTrue会联网下载权重第一次跑如果网络慢建议提前手动下载 vgg16_bn 的权重文件放到~/.cache/torch/hub/checkpoints/否则等待时间可能让你以为程序卡死了。3.2 搭建 DecoderMaxUnpool2d 与卷积层恢复分辨率Decoder 的结构和 Encoder 完全对称上采样用的是MaxUnpool2d然后把特征图通过两个卷积块逐步恢复通道数和分辨率。class SegNetDecoder(nn.Module): def __init__(self, num_classes21): super().__init__() # 每组上采样后接两个卷积层通道数与 Encoder 对称 self.unpool5 nn.MaxUnpool2d(kernel_size2, stride2) self.block5d nn.Sequential( nn.Conv2d(512, 512, kernel_size3, padding1), nn.BatchNorm2d(512), nn.ReLU(inplaceTrue), nn.Conv2d(512, 512, kernel_size3, padding1), nn.BatchNorm2d(512), nn.ReLU(inplaceTrue), ) self.unpool4 nn.MaxUnpool2d(kernel_size2, stride2) self.block4d nn.Sequential( nn.Conv2d(512, 512, kernel_size3, padding1), nn.BatchNorm2d(512), nn.ReLU(inplaceTrue), nn.Conv2d(512, 256, kernel_size3, padding1), nn.BatchNorm2d(256), nn.ReLU(inplaceTrue), ) self.unpool3 nn.MaxUnpool2d(kernel_size2, stride2) self.block3d nn.Sequential( nn.Conv2d(256, 256, kernel_size3, padding1), nn.BatchNorm2d(256), nn.ReLU(inplaceTrue), nn.Conv2d(256, 128, kernel_size3, padding1), nn.BatchNorm2d(128), nn.ReLU(inplaceTrue), ) self.unpool2 nn.MaxUnpool2d(kernel_size2, stride2) self.block2d nn.Sequential( nn.Conv2d(128, 128, kernel_size3, padding1), nn.BatchNorm2d(128), nn.ReLU(inplaceTrue), nn.Conv2d(128, 64, kernel_size3, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), ) self.unpool1 nn.MaxUnpool2d(kernel_size2, stride2) self.block1d nn.Sequential( nn.Conv2d(64, 64, kernel_size3, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.Conv2d(64, num_classes, kernel_size3, padding1), ) def forward(self, x, indices): x self.unpool5(x, indices[4]) x self.block5d(x) x self.unpool4(x, indices[3]) x self.block4d(x) x self.unpool3(x, indices[2]) x self.block3d(x) x self.unpool2(x, indices[1]) x self.block2d(x) x self.unpool1(x, indices[0]) x self.block1d(x) return x这里的关键是forward接收两个参数特征图和 Encoder 保存的 indices 列表。indices[4]对应 Encoder 的pool5indices[3]对应pool4顺序不能乱。最后block1d的第二个卷积层输出通道数是num_classes没有接 ReLU 和 BN因为后面要直接算交叉熵损失。注意MaxUnpool2d的kernel_size和stride必须和对应MaxPool2d完全一致之前踩过坑——Encoder 里kernel_size2, stride2Decoder 写成kernel_size2, stride1结果训练时损失降低到某个点就再也不动了但没有任何报错。这种错误特别隐蔽建议在模型构建后打印一下输出的形状确认是否和输入一致。3.3 损失函数选 CrossEntropyLoss 的边界ignore_index 与类别权重分割任务本质上是一个逐像素分类问题所以最常用的损失函数就是torch.nn.CrossEntropyLoss。它对每个像素位置计算 softmax 交叉熵然后取平均。criterion nn.CrossEntropyLoss(ignore_index255)ignore_index255这个参数很关键。数据集的 mask 中有一些像素是被标注为“未知区域”的在 VOC 数据集中它们的灰度值是 255。如果不忽略这些像素模型会被迫去预测一个根本不存在的类别训练过程会被带偏。另一个参数是weight如果你发现某些类别特别少比如在道路分割中“摩托车”类别占比不到 1%可以给这个类别更高的权重。常见做法是统计训练集中每个类别的像素占比然后取倒数作为权重。但注意如果你的数据集本身不大权重设置过于激进会放大噪声我的经验是weight只在类别占比差距超过一个数量级时才需要设置否则默认的等权重即可。CrossEntropyLoss对输入 logits 和 mask 的尺寸有要求logits 的形状是[N, C, H, W]mask 的形状是[N, H, W]而且 mask 的值必须是 0 到 C-1 之间的整数。如果你的 mask 是 RGB 彩色图需要先转换成单通道的类别索引如果是灰度图但值域是 0 到 255也要先确认最大值是否超过 C-1。这一步出错不会报错但训练出来的结果会是一团糟——损失看似下降但预测输出全是乱的到时候你根本不知道是哪一步出了问题。4. 数据准备与训练从 VOC 到可跑的 Python 工程4.1 数据集组织方式文件夹结构、类别映射与 mask 读取大作业里最常用的公开数据集是 PASCAL VOC 2012它有 21 个类别20 个物体类加 1 个背景图像尺寸不固定。如果你不打算用 VOC也可以用 CamVid道路场景11 类两个数据集的预处理方式几乎一样只是类别映射表不同。这里给出 VOC 数据集的目录组织方式和转换脚本。下载好的 VOC 原始目录长这样VOCdevkit/ └── VOC2012/ ├── JPEGImages/ # 原始图片 ├── SegmentationClass/ # 分割标签PNG 格式调色板模式 └── ImageSets/ └── Segmentation/ ├── train.txt # 训练集图片文件名 └── val.txt # 验证集图片文件名标签文件SegmentationClass里的 PNG 是调色板模式P 模式不是普通的三通道 RGB所以不能用cv2.imread直接读否则读出来的数组是乱的。正确做法是用 PIL 读取并保持P模式from PIL import Image import numpy as np mask Image.open(VOC2012/SegmentationClass/2007_000033.png) # 调色板模式下像素值直接就是类别索引 mask_np np.array(mask, dtypenp.uint8) print(np.unique(mask_np)) # 应该是 0~20 以及 255这里有个容易翻车的细节PIL 打开 P 模式的图像后如果直接.convert(RGB)再转数组得到的是 RGB 三通道值和类别索引完全对不上。必须保持 P 模式读像素值才是类别索引。4.2 自定义 DatasetPIL 读图、调色板转 class id、DataLoader 参数PyTorch 的torchvision.datasets没有现成的 VOC 分割 Dataset 可以直接用于 SegNet 训练需要自己写一个Dataset类。这个类负责按文件列表读取图片和 mask做同步缩放和数据增强。import torch from torch.utils.data import Dataset from PIL import Image import torchvision.transforms as T class VOCSegDataset(Dataset): def __init__(self, img_dir, mask_dir, file_list, img_size(320, 320)): self.img_dir img_dir self.mask_dir mask_dir self.img_size img_size with open(file_list, r) as f: self.names [line.strip() for line in f.readlines()] # 训练时用到的增强随机水平翻转 随机缩放裁剪 self.train_transform T.Compose([ T.RandomHorizontalFlip(p0.5), T.Resize((img_size[0] 20, img_size[1] 20)), T.RandomCrop(img_size), T.ToTensor(), T.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) def __len__(self): return len(self.names) def __getitem__(self, idx): name self.names[idx] img_path f{self.img_dir}/{name}.jpg mask_path f{self.mask_dir}/{name}.png image Image.open(img_path).convert(RGB) mask Image.open(mask_path) # 保持 P 模式 # 图像和 mask 必须使用相同的随机种子才能保证同步增强 seed torch.randint(0, 2**32, (1,)).item() torch.manual_seed(seed) image self.train_transform(image) torch.manual_seed(seed) mask T.Resize(self.img_size, interpolationT.InterpolationMode.NEAREST)(mask) mask torch.as_tensor(np.array(mask), dtypetorch.long) return image, mask代码里的关键点图像和 mask 做随机翻转和裁剪时必须用相同随机种子否则二者会错位。另外mask 的 Resize 插值必须是NEAREST不能是双线性因为双线性插值会在类别边界上产生非整数、不在类别集合内的值等于造出了新类别。__getitem__返回的 mask 是 LongTensor形状是[H, W]没有通道维。这个要在collate_fn里保持DataLoader默认的 collate 能处理但如果你加了别的处理很容易把维度弄丢。from torch.utils.data import DataLoader train_loader DataLoader( VOCSegDataset(VOC2012/JPEGImages, VOC2012/SegmentationClass, train.txt), batch_size8, shuffleTrue, num_workers4, pin_memoryTrue )num_workers4在 Windows 上有时会因为多进程启动方式报错如果你用的是 Windows建议改成num_workers0Linux 上可以保持 4 或者更高。pin_memoryTrue在 GPU 训练时能减少数据拷贝时间但只有在使用 CUDA 时才有意义。4.3 训练循环单卡配置、学习率策略与日志保存训练循环要自己写是意料之中的事PyTorch 没有现成的model.fit。这里给一个标准训练脚本包含学习率衰减、模型保存和每轮验证。import torch.optim as optim from torch.optim.lr_scheduler import StepLR device torch.device(cuda if torch.cuda.is_available() else cpu) model SegNet(num_classes21).to(device) # 注意vgg16_bn 的 pretrained 权重在 model 内部已经加载过了 # 优化器只更新模型的全部参数但可以给 decoder 设置更大的学习率 optimizer optim.SGD([ {params: model.encoder.parameters(), lr: 1e-3}, {params: model.decoder.parameters(), lr: 1e-2}, ], momentum0.9, weight_decay5e-4) scheduler StepLR(optimizer, step_size30, gamma0.1) criterion nn.CrossEntropyLoss(ignore_index255) for epoch in range(50): model.train() total_loss 0.0 for images, masks in train_loader: images images.to(device) masks masks.to(device) optimizer.zero_grad() outputs model(images) # [N, 21, H, W] loss criterion(outputs, masks) loss.backward() optimizer.step() total_loss loss.item() scheduler.step() print(fEpoch {epoch1}, Loss: {total_loss / len(train_loader):.4f}) # 每 5 轮存一次 checkpoint if (epoch 1) % 5 0: torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), }, fcheckpoints/segnet_epoch_{epoch1}.pth)这里的学习率设置有个讲究Encoder 用的是预训练权重学习率给低一些1e-3Decoder 是从头训练学习率给高一些1e-2这种分层学习率是大作业拿高分的小技巧之一。StepLR每 30 轮把学习率乘 0.1如果数据量小、训练轮次少可以改成step_size10。训练时如果发现 loss 不降先确认两件事一是数据增强是不是过度了比如RandomCrop的尺寸和原图相差太大二是CrossEntropyLoss的ignore_index是否正确。这两个都是大作业里最容易出现的问题而且都不会报错只能通过可视化来排查。5. SegNet 图像分割常见问题排查五个踩坑记录5.1 坑 1PIL 读 mask 变成三通道类别数对不上现象训练时模型能跑但 loss 一直在 4 左右降不下去打印np.unique(mask)发现里面有 0 到 255 之外的数字或者类别数突然变成了 200 多。原因PIL 在读取调色板模式的 PNG 时如果代码里写了.convert(RGB)像素值就从类别索引变成了 RGB 值比如红色可能是[255, 0, 0]这样种每个 RGB 组合都成了“新类别”。解决去掉.convert(RGB)保持 P 模式读 mask如果原始 mask 是单通道灰度 PNG确保值域在类别数之内。建议在__getitem__里加一行断言assert mask_np.max() num_classes or mask_np.max() 255这样即使出问题也能第一时间发现。5.2 坑 2MaxUnpool2d 报错 size mismatch现象前向传播时报错提示Size mismatch具体信息是got input size [N, C, H, W], but expected size [N, C, 2H, 2W]。原因MaxUnpool2d会根据输入特征图和索引推出期望的输出尺寸。如果 Encoder 里池化的kernel_size和 Decoder 里反池化的不一致或者输入图像的尺寸不是 2 的幂次索引就会和特征图对不上。解决保持kernel_size和stride完全一致输入图像的宽高必须是 2 的幂次比如 320、384 都可以。如果你用的是不规则的尺寸建议在训练前统一 resize 到 320×320 或 384×384不要用-1之类的自适应尺寸。这个坑在写论文用的自建数据集上尤其常见因为很多人直接拿原始尺寸训练。5.3 坑 3mIoU 一直很低原来是损失反向传播到了预训练权重现象训练了 20 轮训练 loss 正常下降但验证集 mIoU 只有不到 0.2且不随训练上升。原因一种可能是没有冻结 Encoder 的 BatchNorm 层。加载的 VGG16 预训练权重带有 BN 层的running_mean和running_var如果数据集的分布和 ImageNet 差异很大比如医学图像BN 层在微调中会不断更新这些统计量导致特征分布不稳定。解决在训练前冻结所有 BN 层的参数for module in model.encoder.modules(): if isinstance(module, nn.BatchNorm2d): module.eval() # 固定 running_mean 和 running_var注意要在model.train()之后设置module.eval()这个操作只影响 BN 的统计量更新不影响其他层的训练。另一种做法是换成GroupNorm但这里不展开。5.4 坑 4显存溢出切 patch 后边缘出现黑色条带现象训练一张 1024×1024 的大图直接 OOM于是切成 256×256 的 patch 训练但预测时把 patch 拼回原图边缘出现黑色条带。原因切 patch 训练时模型只见过 256×256 的输入感受野有限。预测时如果用全图输入模型的表现和训练时不一致如果使用 patch 拼接patch 之间有重叠部分但重叠区域没有做平滑融合边界处概率不一致取 argmax 后就会出现明显的条带。解决推理时采用滑窗加重叠平均。每个位置至少被预测 2 次取概率平均后再 argmax。重叠比例取 1/4 宽度即可效果明显。如果不想写也可以直接把输入 resize 到训练尺寸再预测但这样会损失小目标的精度。5.5 坑 5验证时预测图全黑softmax 维度取错现象验证阶段把outputs传入torch.argmax得到的结果保存成图片后整张图是全黑的或者只有零星几个点。原因outputs的形状是[N, C, H, W]argmax应该在维度 1通道维上取但很多人习惯性写成torch.argmax(outputs, dim0)这样得到的是对 batch 维取最大值索引而不是对类别维取最大值。解决统一写成pred torch.argmax(outputs, dim1) # [N, H, W]然后保存的时候注意索引 0 对应的颜色是背景。如果你用的调色板是从 0 开始的保存为 PNG 时要用P模式和自定义调色板否则像 PIL 默认保存的 RGB 模式会把 0~20 的索引当成灰度看起来就是黑糊糊的。6. 用可视化验证训练效果质量图、IoU 分析与 ONNX 导出训练结束后你不能只上交一个.pth权重文件就算完事。大作业要想拿高分必须有一套让人看得懂的验证结果输出。这里给一个最简单也最有效的验证三件套损失曲线、逐类 IoU 柱状图、和原图像素级的可视化对比。先写一个脚本一次性输出所有验证集图片的预测结果并合成一张对比图import matplotlib.pyplot as plt # 类别颜色表VOC 21 类对应的调色板 palette [ (0, 0, 0), (128, 0, 0), (0, 128, 0), (128, 128, 0), (0, 0, 128), (128, 0, 128), (0, 128, 128), (128, 128, 128), (64, 0, 0), (192, 0, 0), (64, 128, 0), (192, 128, 0), (64, 0, 128), (192, 0, 128), (64, 128, 128), (192, 128, 128), (0, 64, 0), (128, 64, 0), (0, 192, 0), (128, 192, 0), (0, 64, 128) ] def colorize_mask(mask_np, palette): h, w mask_np.shape color_mask np.zeros((h, w, 3), dtypenp.uint8) for i in range(len(palette)): color_mask[mask_np i] palette[i] return color_mask model.eval() with torch.no_grad(): for images, masks in val_loader: images images.to(device) outputs model(images) pred torch.argmax(outputs, dim1).cpu().numpy() # 取 batch 中第一张图做可视化 img images[0].cpu().numpy().transpose(1, 2, 0) img img * np.array([0.229, 0.224, 0.225]) np.array([0.485, 0.456, 0.406]) img (img * 255).astype(np.uint8) pred_color colorize_mask(pred[0], palette) gt_color colorize_mask(masks[0].cpu().numpy(), palette) fig, axes plt.subplots(1, 3, figsize(15, 5)) axes[0].imshow(img); axes[0].set_title(Input) axes[1].imshow(gt_color); axes[1].set_title(Ground Truth) axes[2].imshow(pred_color); axes[2].set_title(Prediction) plt.savefig(fvis/result_{idx}.png, bbox_inchestight) plt.close()这段脚本里用colorize_mask把类别索引映射成彩色图视觉上要比灰度图直观得多。注意图片反标准化时必须先乘标准差再加均值顺序反了颜色会失真。IoU 分析方面不要只报一个平均 mIoU建议对每个类单独算 IoU然后用柱状图展示。大作业报告里贴一张逐类 IoU 柱状图再配一段文字说明边界类比如瓶子、椅子为什么难分割这个深度就已经超出很多同水平作业了。如果还想进一步加分可以把训练好的模型用torch.onnx.export导出为 ONNX然后在报告里附一句“模型已转为 ONNX 格式方便后续部署到移动端或服务端”。但要注意 SegNet 的MaxUnpool2d在导出 ONNX 时的算子兼容性不同版本的 PyTorch 对MaxUnpool的支持不完全一致如果导出报错可以用torch.jit.trace作为替代方案。这个不是必需步骤但确实能让答辩老师眼前一亮。最后一个建议来自我的血泪经验训练完成后把模型在 3 张不同场景的测试图上的预测结果保存下来一张是物体密集的、一张是大面积背景的、一张是光照异常的。然后把这三张图放进报告里分别分析模型的表现。这样的分析比一张总图更有说服力也能体现出你对自己模型边界的理解。图像分割实验的“黑匣子”程度非常高训练过程里很多问题都要靠可视化去诊断。养成每次训练后都保存可视化结果的习惯比盯着 loss 数字有用得多。我在第一次跑 SegNet 的时候就是因为没有保存中间结果翻车了好几次才意识到这个问题。希望这些经验能帮你少走点弯路祝你好运。本文还有配套的精品资源点击获取