简介基于U-Net的道路目标语义分割系统面向具备一定PyTorch基础的计算机视觉开发者解决道路场景下车辆、行人、建筑物等目标的像素级分类问题可服务于自动驾驶、智能交通及城市规划等应用场景。压缩包共7个文件以5个Python源码为主体分别承担数据预处理、模型构建、训练循环、预测推理与工具封装等职责另含1个环境依赖配置和1个说明文档便于快速还原运行环境。整体包体仅10KB结构紧凑适合直接对照学习。目前已有118人学习下载。项目采用U-Net经典的编码器-解码器架构能有效保留空间语义信息代码模块划分清晰从数据加载到模型评估形成完整闭环可帮助开发者理解语义分割的核心流程同时提供调参与扩展思路便于迁移至其他图像分割任务。1. 道路目标语义分割为什么我最终回到U-Net第一次接到道路目标分割的需求时我第一反应是直接找一个带预训练权重的分割模型比如 DeepLabV3 或 PSPNet毕竟这类模型在公开榜单上表现漂亮。但真正落到自己的数据上我发现两个问题一是道路场景类别杂道路、车辆、行人、骑行者、交通标志、路缘石都得分开并不是简单的前景背景二分类二是车载算力有限模型不能太折腾。兜了一圈后我最终还是用 U-Net 搭了一套可复现的 PyTorch 方案。U-Net 的编码器-解码器结构和跳跃连接在中等尺寸数据集上能把细节边界还原得相当好训练成本也可控。这套系统适合那些手里有少量自己标注数据、想把道路目标分割快速跑起来的从业者。下面这篇笔记就是我把整个系统从数据到推理完整拆一遍的过程包含可直接抄的代码和四个最容易翻车的坑。2. 把 U-Net 搭起来PyTorch 实现与跳跃连接的真正意义2.1 编码器用卷积块和下采样压缩空间信息U-Net 的编码器本质是一串“卷积 池化”的堆叠。每个阶段先做两次卷积再通过最大池化把特征图缩小一半。这样做有两个目的一是扩大感受野让网络在高层能看到更大范围的上下文比如判断某个区域到底是道路还是人行道二是降低计算量高层特征通道数翻倍空间尺寸减半参数量增长可控。我习惯先把一个双卷积块抽出来后续编码和解码都能复用。下面这个块用卷积 BatchNorm ReLU 的标准组合卷积核统一用 3×3padding 为 1保证特征图尺寸不变。import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() self.conv nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size3, padding1), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue), nn.Conv2d(out_channels, out_channels, kernel_size3, padding1), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.conv(x)这里有几个参数值得注意。in_channels第一层是 3对应 RGB 输入out_channels从 64 开始每下一层翻倍。BatchNorm2d在 batch size 较小时会抖动明显如果显存只允许 batch size 为 2 或 4建议先不开 BatchNorm改成 GroupNorm 或 InstanceNorm否则训练曲线会像心电图一样上下跳。inplaceTrue是为了省显存反向传播时 ReLU 不需要保留输入能省一点是一点。编码器部分就是串行调用这些块和池化层。我用 MaxPool2d 而不是 Conv2d stride2 做下采样原因是池化没有可学习参数训练更稳定且不容易丢失低层位置信息。2.2 解码器上采样与跳跃连接编码器把一张 512×512 的图逐步压缩到 32×32 的高维特征但分割要求输出和原图一样大小的逐像素类别所以需要一个解码器把空间分辨率逐层恢复。解码器的核心操作是转置卷积上采样再把对应编码器层的特征图拼回来。这个“拼回来”的操作就是 U-Net 的跳跃连接。它解决的是一个很实际的问题经过四次池化后高层特征虽然语义信息丰富但边界和细节已经被抹掉了。跳跃连接直接把浅层的细节特征传给解码器让网络在恢复空间尺寸的同时还有足够的信息去修正边缘。class UNet(nn.Module): def __init__(self, n_classes8): super().__init__() self.enc1 DoubleConv(3, 64) self.pool1 nn.MaxPool2d(2) self.enc2 DoubleConv(64, 128) self.pool2 nn.MaxPool2d(2) self.enc3 DoubleConv(128, 256) self.pool3 nn.MaxPool2d(2) self.enc4 DoubleConv(256, 512) self.pool4 nn.MaxPool2d(2) self.bridge DoubleConv(512, 1024) self.up1 nn.ConvTranspose2d(1024, 512, kernel_size2, stride2) self.dec1 DoubleConv(1024, 512) self.up2 nn.ConvTranspose2d(512, 256, kernel_size2, stride2) self.dec2 DoubleConv(512, 256) self.up3 nn.ConvTranspose2d(256, 128, kernel_size2, stride2) self.dec3 DoubleConv(256, 128) self.up4 nn.ConvTranspose2d(128, 64, kernel_size2, stride2) self.dec4 DoubleConv(128, 64) self.out_conv nn.Conv2d(64, n_classes, kernel_size1) def forward(self, x): e1 self.enc1(x) e2 self.enc2(self.pool1(e1)) e3 self.enc3(self.pool2(e2)) e4 self.enc4(self.pool3(e3)) b self.bridge(self.pool4(e4)) d1 self.dec1(torch.cat([self.up1(b), e4], dim1)) d2 self.dec2(torch.cat([self.up2(d1), e3], dim1)) d3 self.dec3(torch.cat([self.up3(d2), e2], dim1)) d4 self.dec4(torch.cat([self.up4(d3), e1], dim1)) return self.out_conv(d4)torch.cat的dim1是在通道维上拼接。注意上采样后的特征图尺寸必须和对应编码器层完全一致否则拼接会直接报错。如果输入尺寸不是 2 的整数次幂比如 512×512 没问题但如果是 500×500四次池化后会出现 31×31 和 32×32 对不上的情况。我一般会把所有训练图统一 resize 成 512×512 或 384×384避免这类对齐问题。最后一个 1×1 卷积把 64 通道映射到类别数n_classes。这里输出的是未经过 softmax 的 logits训练时配合交叉熵损失使用推理时再取 argmax效率更高。2.3 输出头与损失函数分割任务里输出头很少单独设计但损失函数值得多说一句。最常见的是nn.CrossEntropyLossPyTorch 里它对logits直接计算不需要手动做 softmax。对于道路目标这类场景类别之间数量极不平衡道路可能占图像面积的 60%行人可能只占 0.5%。直接用交叉熵网络会倾向于把所有像素都预测成道路mIoU 看起来还行行人这一类直接废掉。所以我常把交叉熵和 Dice Loss 结合起来。Dice Loss 对前景小目标更友好计算的是预测区域和真实区域的像素重叠程度。简单实现如下def dice_loss(pred, target, smooth1.0): pred torch.softmax(pred, dim1) target_onehot torch.nn.functional.one_hot(target, num_classespred.shape[1]).permute(0, 3, 1, 2).float() intersection (pred * target_onehot).sum(dim(2, 3)) union pred.sum(dim(2, 3)) target_onehot.sum(dim(2, 3)) dice (2.0 * intersection smooth) / (union smooth) return 1.0 - dice.mean()组合方式可以是total_loss ce_loss 0.5 * dice_loss。这里smooth参数设为 1.0 是防分母为 0 的常规操作。one_hot需要 target 是长整型索引图不能是 RGB 标签图这个细节在下一章专门讲。3. 数据准备与增强道路场景语义分割的关键一步3.1 数据集目录结构与标签处理很多第一次做分割的人会直接在原始 JPEG 图上跑标签图拿出来却是一张彩色 PNG然后报错说通道对不上。语义分割的标签必须是单通道索引图每个像素的数值代表一个类别编号。比如 0 代表道路1 代表车辆2 代表行人。如果你手里是 Cityscapes 那种彩色标签图必须进行颜色映射转换。我一般把数据按下面这种方式组织dataset/ images/ 000001.jpg 000002.jpg labels/ 000001.png 000002.png图像和标签文件名保持一致代码里通过文件名前缀来关联。标签用 PNG 格式保存无损索引图不要用 JPG因为 JPG 压缩会在标签边缘产生噪点导致类别编号被改动。很多时候你拿到的标注是 RGB 颜色图比如道路是紫色、车辆是蓝色。要转成索引图必须维护一个颜色到类别编号的映射表。下面这段代码是我常用的转换方式color2id { (128, 64, 128): 0, # 道路 ( 70, 70, 70): 1, # 建筑物 (250,170, 30): 2, # 车辆 (220, 20, 60): 3, # 行人 } def rgb_to_index(rgb_label): h, w, _ rgb_label.shape index np.zeros((h, w), dtypenp.uint8) for color, class_id in color2id.items(): mask (rgb_label np.array(color)).all(axis-1) index[mask] class_id return index这段代码看起来简单但实际很容易踩坑。第一个坑是 OpenCV 读图默认 BGR 顺序cv2.imread读进来后rgb_label的通道顺序是反的你直接比对 (128, 64, 128) 永远匹配不上。第二个坑是颜色值可能有 JPEG 压缩带来的微小偏移比如 128 变成 127导致掩码全部漏掉。我一般要求标注导出时必须是 PNG 无损格式且不要用 JPEG 保存标注图。如果你发现某些类别的像素数量明显异常少先从这两点排查。3.2 数据增强不要破坏道路结构的空间关系图像分类里常用的随机裁剪、翻转和旋转在分割任务里要更谨慎。比如随机旋转 90 度对道路场景来说可能把“上方是天空、下方是路面”的空间关系打乱模型学到的先验知识被破坏。我常用的增强策略是水平翻转、小幅旋转±5 度以内、随机缩放0.8 到 1.2 倍、随机亮度和对比度调整。这里我直接用albumentations它对图像和掩码同步变换不用手工保证两套变换一致import albumentations as A train_transform A.Compose([ A.Resize(height512, width512), A.HorizontalFlip(p0.5), A.RandomScale(scale_limit(0.8, 1.2), p0.4), A.RandomBrightnessContrast(brightness_limit0.1, contrast_limit0.1, p0.4), ])注意每一项的p参数是应用概率不要全部设为 1。RandomScale会改变图像和掩码的尺寸如果后面接Resize顺序很重要。我习惯把Resize放在最前面保证后续增强在固定尺寸下进行否则RandomScale之后再Resize等于缩放被抵消了增强效果不明显。对于道路分割我不推荐使用粗粒度的 Cutout 或随机遮挡。因为车辆和行人这类目标本身就很依赖完整轮廓中间挖掉一块模型很可能就把那个区域误判成路面。如果确实需要正则化可以试试在特征图上做 Dropout而不是在输入图像上做剪贴。3.3 自定义 Dataset 与数据加载器PyTorch 里写一个自定义 Dataset 是固定套路。我把它拆成几步初始化时读入所有文件路径__getitem__里加载图像和标签做训练增强最后把标签转成torch.long张量。import os import cv2 import torch from torch.utils.data import Dataset class RoadDataset(Dataset): def __init__(self, img_dir, label_dir, transformNone, color_mapNone): self.img_dir img_dir self.label_dir label_dir self.transform transform self.color_map color_map self.img_names sorted(os.listdir(img_dir)) def __len__(self): return len(self.img_names) def __getitem__(self, idx): img_name self.img_names[idx] img_path os.path.join(self.img_dir, img_name) label_path os.path.join(self.label_dir, img_name.replace(.jpg, .png)) image cv2.imread(img_path) image cv2.cvtColor(image, cv2.COLOR_BGR2RGB) label cv2.imread(label_path, cv2.IMREAD_UNCHANGED) if len(label.shape) 3: label self.rgb_to_index(label) if self.transform: transformed self.transform(imageimage, masklabel) image transformed[image] label transformed[mask] image torch.from_numpy(image.transpose(2, 0, 1)).float() / 255.0 label torch.from_numpy(label).long() return image, labelcv2.imread(label_path, cv2.IMREAD_UNCHANGED)这里必须指定UNCHANGED否则 OpenCV 会把单通道 PNG 读成三通道。如果标签是索引图label.shape是二维直接作为 mask 传给 albumentations如果是彩色图需要先转索引。image.transpose(2, 0, 1)是从 HWC 转 CHW标准化通过除以 255.0 实现我一般不在 Dataset 里再做复杂的 Normalize因为这会影响后续的训练稳定性调试。数据加载器里num_workers建议先从 4 开始batch_size视显存定。pin_memoryTrue可以加速 GPU 拷贝但如果你用的是 Windows 系统并且 workers 经常报错可以先关掉试一下。4. 训练与调参从 loss 曲线到 mIoU4.1 优化器与学习率策略U-Net 训练最常见的优化器选择是 Adam 或 SGD。Adam 收敛快一开始 loss 掉得猛但最后 mIoU 可能不如 SGD 磨得高。我的习惯是先用 AdamW 跑 10 个 epoch把网络结构里的坑都排掉比如显存不足、标签错位确认没有问题后再用带 momentum 的 SGD 重训一遍通常能再多涨 2 到 3 个点。学习率设置上ResNet 系主干一般用 1e-4 起步U-Net 我推荐 2e-4 到 5e-4。配合 poly 学习率衰减策略也就是每个 iteration 都降低学习率公式为lr base_lr * (1 - iter / total_iter) ** 0.9。这种方式比固定步长衰减更适合分割任务后期微调更细腻。from torch.optim import AdamW from torch.optim.lr_scheduler import LambdaLR optimizer AdamW(model.parameters(), lr2e-4, weight_decay1e-4) def poly_scheduler(epoch, warmup_epochs2, total_epochs60, base_lr2e-4): if epoch warmup_epochs: return (epoch 1) / warmup_epochs return (1 - (epoch - warmup_epochs) / (total_epochs - warmup_epochs)) ** 0.9 scheduler LambdaLR(optimizer, lr_lambdalambda epoch: poly_scheduler(epoch))weight_decay我设 1e-4太小起不到约束作用太大会导致模型欠拟合。前两个 epoch 做 warmup让学习率从 0.25 倍慢慢升到基准值这一步非常关键。如果一开始就用大学习率BatchNorm 的 moving mean 还没稳定训练 loss 会剧烈波动甚至直接产生 NaN。4.2 评估指标mIoU 的计算方式与细节分割任务最通用的评估指标是 mIoU计算每个类别的预测集合与真实集合的交并比然后取平均。很多人以为用 PyTorch 里现成的函数就行但实际上没这么简单自己写几行反而更可控。def compute_miou(pred_logits, target): pred torch.argmax(pred_logits, dim1) ious [] for cls_id in range(pred_logits.shape[1]): pred_mask (pred cls_id) target_mask (target cls_id) intersection (pred_mask target_mask).sum().float() union (pred_mask | target_mask).sum().float() if union.item() 0: ious.append(torch.tensor(1.0)) else: ious.append(intersection / union) return torch.stack(ious).mean().item()这里有个常见争议某个类别在当前 batch 中完全没有出现它的 union 是 0mIoU 应该怎么算我见过两种做法一种是跳过这个类别只对出现的类别算平均另一种是直接视为 1.0因为真实和预测都是空区域完全重合。在道路分割场景里如果一个 batch 里没有行人我会把它按 1.0 处理否则验证集指标会被这种“偶然缺类”的 batch 拉低显得模型很差误导调参。4.3 训练循环完整骨架训练循环的核心就三步前向算 loss、反向传播、更新权重。但需要注意model.train()和model.eval()的切换以及验证阶段需要关闭梯度计算否则显存会被验证过程占满。train_loader DataLoader(train_set, batch_size8, shuffleTrue, num_workers4) val_loader DataLoader(val_set, batch_size8, shuffleFalse, num_workers4) model UNet(n_classes5).cuda() criterion nn.CrossEntropyLoss(ignore_index-1) optimizer AdamW(model.parameters(), lr2e-4) scheduler LambdaLR(optimizer, lr_lambdalambda epoch: poly_scheduler(epoch)) best_iou 0.0 for epoch in range(60): model.train() train_loss 0.0 for images, labels in train_loader: images, labels images.cuda(), labels.cuda() optimizer.zero_grad() logits model(images) loss criterion(logits, labels) loss.backward() optimizer.step() train_loss loss.item() model.eval() val_iou 0.0 with torch.no_grad(): for images, labels in val_loader: images, labels images.cuda(), labels.cuda() logits model(images) val_iou compute_miou(logits, labels) val_iou / len(val_loader) scheduler.step() if val_iou best_iou: best_iou val_iou torch.save(model.state_dict(), unet_road_best.pt)ignore_index-1是一个很实用的设置。如果你的标签图里有未标注区域把那些像素的值设成 -1交叉熵会自动跳过这些像素不影响梯度更新。我认为这个参数比手动做 mask 简单得多。保存模型时只保存state_dict()不要整个模型都 pickle 保存否则换环境时容易因为类的路径不匹配而加载失败。加载时先实例化模型再load_state_dict这是一个让人少掉不少头发的习惯。5. 避坑清单道路分割最常见的五个问题排查5.1 标签错位loss 小但预测图整体错位现象训练 loss 降得很顺利验证 mIoU 也不差但把预测图叠到原图上一看整辆车或整条路的边缘都朝一个方向偏移了几像素像是标签和图像没对齐。原因数据预处理里图像 resize 和标签 resize 用了不同的插值方式。比如图像用双线性插值标签用最近邻插值缩放尺度相同时还好但只要图像和标签的尺寸不是严格同步就会错位。另一个常见原因是某些标注工具导出时图像和标签的坐标系原点不一致导致裁剪时左右、上下偏差。解决我用的是 albumentations 同步变换基本能避免这类问题。如果是自写代码一定保证图像和标签共享同一个变换参数矩阵比如cv2.warpAffine时同一个 M 矩阵必须同时作用于图像和标签。标签千万不能用双线性插值因为类别编号是离散的插值会生成不存在的类别比如类别 1 和类别 2 之间插出 1.2。5.2 loss 是 NaN或者头几个 epoch 就崩现象训练刚开始loss 显示 nan或者第一轮还没跑完直接报错CUDA error: device-side assert triggered。原因最常见的是标签值超出了模型输出类别数n_classes。比如模型n_classes5但标签图里有像素值是 6交叉熵计算时越界触发 CUDA assert。第二个原因是 batch size 太小加上 BatchNorm 不稳定梯度爆炸导致参数变成 NaN。解决第一步先统计标签的像素取值集合确认最大值小于n_classesnp.unique(label)看一遍。第二步在网络的forward输出后加一个断言assert output.shape[1] target.max().item()只保留到调试阶段。如果是 BatchNorm 问题把 batch size 提到 8 以上或者把 BatchNorm 换成 GroupNorm通常立竿见影。5.3 道路占绝大多数车辆和行人根本学不出来现象训练结束后单看 mIoU 有 85%但打开预测图所有非道路目标要么消失要么糊成一片。原因这是典型的类别不平衡。道路像素占图中 70% 以上交叉熵损失里道路类贡献的梯度远大于车辆和行人网络只要把道路分对loss 就很低小目标成了可以被牺牲的部分。解决一是调整损失权重给稀少类别更高的权重比如torch.class_weight根据像素频率的逆来设置。二是直接把dice_loss的权重调大因为 Dice 对每个类别单独计算。三是检查数据增强里是否加了过多的随机缩放车辆在缩小后只有十几个像素经过连续池化后信息基本丢失。一个小技巧是把训练图像裁成多个 256×256 patch每个 patch 尽量包含一张图里的小目标让模型更频繁地看到行人这类类别。5.4 推理结果出现棋盘格或密密麻麻的小洞现象模型在训练集上效果不错测试视频里道路中间却出现棋盘格状的小块或者单个车辆目标内部有孔洞。原因棋盘格通常来自转置卷积的叠加尤其是当上采样层数多、通道数大时转置卷积会对某些频率产生过强的响应。道路本身的纹理也很复杂模型对高置信度的道路区域内部做 argmax 时可能因为相邻像素特征不一致而输出不同类别。孔洞则说明分类器对目标内部特征的把握不够把车辆中间几块误判成其他类别。解决上采样替换成nn.Upsample(scale_factor2, modebilinear, align_cornersTrue)再接卷积而不是用转置卷积。虽然理论上转置卷积可学习但在小数据集上容易过拟合棋盘格几乎都是它带来的。推理后处理方面可以加一个简单的多数投票对预测图的局部窗口取众数过滤孤立点。用 OpenCV 的modeFilter或morphologyEx都行但注意不要过度处理否则会把车辆和行人边缘磨平。5.5 显存不够batch size 调小后训练效果断崖式下降现象batch size 从 16 降到 4 之后同样的 epoch 数验证 mIoU 掉了十来个点。原因BatchNorm 在小 batch 下统计均值方差抖动厉害模型每个 step 看到的归一化分布都不一样训练不稳定。很多人误以为减 batch size 只是慢一点实际上对分割模型影响很大因为分割模型的 BatchNorm 层很多。解决如果显存只允许 batch size 4建议把所有BatchNorm2d替换成GroupNorm一般分组数设为 32 或 8效果接近大 batch 下的 BatchNorm。另一个做法是使用梯度累积每 4 个小 batch 累积一次梯度再更新等价于 batch size 翻 4 倍但 BN 的统计仍然是基于小 batch 的所以并不能完全解决问题。最彻底的办法是改用预训练编码器比如用 ResNet34 作为 encoder 的 U-Net 变体显存占用小很多而且收敛更快。6. 推理优化与模型落地从单张图到连续帧6.1 滑动窗口推理与大图切块实际落地时你拿到的图片可能不是 512×512而是 1920×1080 的车载截图。直接 resize 到 512×512 会丢失大量细节小目标直接糊掉。我的做法是滑窗推理把原图切成多个 512×512 的块块之间重叠 64 像素每个块独立预测最后把预测结果拼回原图重叠区域取 argmax 或多数投票。这样可以保留原始分辨率代价是推理次数增加。重叠区域的处理我一般直接取第一个块的预测结果不做融合因为融合逻辑复杂且容易在边界引入新问题。如果你追求更平滑的结果可以对重叠部分按像素距离做加权平均但实测差异不大没必要为此增加代码复杂度。6.2 从 PyTorch 到 TorchScript 导出训练完成后如果要把模型放到 Python 服务或 C 端侧可以用 TorchScript 导出省掉 PyTorch 运行时。导出前务必进入model.eval()模式否则 BatchNorm 的 moving average 不会被固定导出的模型结果会不稳定。model UNet(n_classes5) model.load_state_dict(torch.load(unet_road_best.pt)) model.eval() dummy torch.randn(1, 3, 512, 512) traced torch.jit.trace(model, dummy) traced.save(unet_road.pt)导出后用traced(images)验证一次和原模型逐像素对比误差应该在1e-5以下。如果误差很大大概率是torch.jit.trace路径里有条件分支或动态循环U-Net 的静态卷积结构一般不会有这个问题。从那以后我每次训练完都会强制走一遍这个导出流程确保交付出去的模型不是只在训练代码里能跑的“黑匣子”。希望帮到你。本文还有配套的精品资源点击获取
U-Net实战：基于PyTorch的道路目标语义分割与避坑指南
简介基于U-Net的道路目标语义分割系统面向具备一定PyTorch基础的计算机视觉开发者解决道路场景下车辆、行人、建筑物等目标的像素级分类问题可服务于自动驾驶、智能交通及城市规划等应用场景。压缩包共7个文件以5个Python源码为主体分别承担数据预处理、模型构建、训练循环、预测推理与工具封装等职责另含1个环境依赖配置和1个说明文档便于快速还原运行环境。整体包体仅10KB结构紧凑适合直接对照学习。目前已有118人学习下载。项目采用U-Net经典的编码器-解码器架构能有效保留空间语义信息代码模块划分清晰从数据加载到模型评估形成完整闭环可帮助开发者理解语义分割的核心流程同时提供调参与扩展思路便于迁移至其他图像分割任务。1. 道路目标语义分割为什么我最终回到U-Net第一次接到道路目标分割的需求时我第一反应是直接找一个带预训练权重的分割模型比如 DeepLabV3 或 PSPNet毕竟这类模型在公开榜单上表现漂亮。但真正落到自己的数据上我发现两个问题一是道路场景类别杂道路、车辆、行人、骑行者、交通标志、路缘石都得分开并不是简单的前景背景二分类二是车载算力有限模型不能太折腾。兜了一圈后我最终还是用 U-Net 搭了一套可复现的 PyTorch 方案。U-Net 的编码器-解码器结构和跳跃连接在中等尺寸数据集上能把细节边界还原得相当好训练成本也可控。这套系统适合那些手里有少量自己标注数据、想把道路目标分割快速跑起来的从业者。下面这篇笔记就是我把整个系统从数据到推理完整拆一遍的过程包含可直接抄的代码和四个最容易翻车的坑。2. 把 U-Net 搭起来PyTorch 实现与跳跃连接的真正意义2.1 编码器用卷积块和下采样压缩空间信息U-Net 的编码器本质是一串“卷积 池化”的堆叠。每个阶段先做两次卷积再通过最大池化把特征图缩小一半。这样做有两个目的一是扩大感受野让网络在高层能看到更大范围的上下文比如判断某个区域到底是道路还是人行道二是降低计算量高层特征通道数翻倍空间尺寸减半参数量增长可控。我习惯先把一个双卷积块抽出来后续编码和解码都能复用。下面这个块用卷积 BatchNorm ReLU 的标准组合卷积核统一用 3×3padding 为 1保证特征图尺寸不变。import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() self.conv nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size3, padding1), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue), nn.Conv2d(out_channels, out_channels, kernel_size3, padding1), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.conv(x)这里有几个参数值得注意。in_channels第一层是 3对应 RGB 输入out_channels从 64 开始每下一层翻倍。BatchNorm2d在 batch size 较小时会抖动明显如果显存只允许 batch size 为 2 或 4建议先不开 BatchNorm改成 GroupNorm 或 InstanceNorm否则训练曲线会像心电图一样上下跳。inplaceTrue是为了省显存反向传播时 ReLU 不需要保留输入能省一点是一点。编码器部分就是串行调用这些块和池化层。我用 MaxPool2d 而不是 Conv2d stride2 做下采样原因是池化没有可学习参数训练更稳定且不容易丢失低层位置信息。2.2 解码器上采样与跳跃连接编码器把一张 512×512 的图逐步压缩到 32×32 的高维特征但分割要求输出和原图一样大小的逐像素类别所以需要一个解码器把空间分辨率逐层恢复。解码器的核心操作是转置卷积上采样再把对应编码器层的特征图拼回来。这个“拼回来”的操作就是 U-Net 的跳跃连接。它解决的是一个很实际的问题经过四次池化后高层特征虽然语义信息丰富但边界和细节已经被抹掉了。跳跃连接直接把浅层的细节特征传给解码器让网络在恢复空间尺寸的同时还有足够的信息去修正边缘。class UNet(nn.Module): def __init__(self, n_classes8): super().__init__() self.enc1 DoubleConv(3, 64) self.pool1 nn.MaxPool2d(2) self.enc2 DoubleConv(64, 128) self.pool2 nn.MaxPool2d(2) self.enc3 DoubleConv(128, 256) self.pool3 nn.MaxPool2d(2) self.enc4 DoubleConv(256, 512) self.pool4 nn.MaxPool2d(2) self.bridge DoubleConv(512, 1024) self.up1 nn.ConvTranspose2d(1024, 512, kernel_size2, stride2) self.dec1 DoubleConv(1024, 512) self.up2 nn.ConvTranspose2d(512, 256, kernel_size2, stride2) self.dec2 DoubleConv(512, 256) self.up3 nn.ConvTranspose2d(256, 128, kernel_size2, stride2) self.dec3 DoubleConv(256, 128) self.up4 nn.ConvTranspose2d(128, 64, kernel_size2, stride2) self.dec4 DoubleConv(128, 64) self.out_conv nn.Conv2d(64, n_classes, kernel_size1) def forward(self, x): e1 self.enc1(x) e2 self.enc2(self.pool1(e1)) e3 self.enc3(self.pool2(e2)) e4 self.enc4(self.pool3(e3)) b self.bridge(self.pool4(e4)) d1 self.dec1(torch.cat([self.up1(b), e4], dim1)) d2 self.dec2(torch.cat([self.up2(d1), e3], dim1)) d3 self.dec3(torch.cat([self.up3(d2), e2], dim1)) d4 self.dec4(torch.cat([self.up4(d3), e1], dim1)) return self.out_conv(d4)torch.cat的dim1是在通道维上拼接。注意上采样后的特征图尺寸必须和对应编码器层完全一致否则拼接会直接报错。如果输入尺寸不是 2 的整数次幂比如 512×512 没问题但如果是 500×500四次池化后会出现 31×31 和 32×32 对不上的情况。我一般会把所有训练图统一 resize 成 512×512 或 384×384避免这类对齐问题。最后一个 1×1 卷积把 64 通道映射到类别数n_classes。这里输出的是未经过 softmax 的 logits训练时配合交叉熵损失使用推理时再取 argmax效率更高。2.3 输出头与损失函数分割任务里输出头很少单独设计但损失函数值得多说一句。最常见的是nn.CrossEntropyLossPyTorch 里它对logits直接计算不需要手动做 softmax。对于道路目标这类场景类别之间数量极不平衡道路可能占图像面积的 60%行人可能只占 0.5%。直接用交叉熵网络会倾向于把所有像素都预测成道路mIoU 看起来还行行人这一类直接废掉。所以我常把交叉熵和 Dice Loss 结合起来。Dice Loss 对前景小目标更友好计算的是预测区域和真实区域的像素重叠程度。简单实现如下def dice_loss(pred, target, smooth1.0): pred torch.softmax(pred, dim1) target_onehot torch.nn.functional.one_hot(target, num_classespred.shape[1]).permute(0, 3, 1, 2).float() intersection (pred * target_onehot).sum(dim(2, 3)) union pred.sum(dim(2, 3)) target_onehot.sum(dim(2, 3)) dice (2.0 * intersection smooth) / (union smooth) return 1.0 - dice.mean()组合方式可以是total_loss ce_loss 0.5 * dice_loss。这里smooth参数设为 1.0 是防分母为 0 的常规操作。one_hot需要 target 是长整型索引图不能是 RGB 标签图这个细节在下一章专门讲。3. 数据准备与增强道路场景语义分割的关键一步3.1 数据集目录结构与标签处理很多第一次做分割的人会直接在原始 JPEG 图上跑标签图拿出来却是一张彩色 PNG然后报错说通道对不上。语义分割的标签必须是单通道索引图每个像素的数值代表一个类别编号。比如 0 代表道路1 代表车辆2 代表行人。如果你手里是 Cityscapes 那种彩色标签图必须进行颜色映射转换。我一般把数据按下面这种方式组织dataset/ images/ 000001.jpg 000002.jpg labels/ 000001.png 000002.png图像和标签文件名保持一致代码里通过文件名前缀来关联。标签用 PNG 格式保存无损索引图不要用 JPG因为 JPG 压缩会在标签边缘产生噪点导致类别编号被改动。很多时候你拿到的标注是 RGB 颜色图比如道路是紫色、车辆是蓝色。要转成索引图必须维护一个颜色到类别编号的映射表。下面这段代码是我常用的转换方式color2id { (128, 64, 128): 0, # 道路 ( 70, 70, 70): 1, # 建筑物 (250,170, 30): 2, # 车辆 (220, 20, 60): 3, # 行人 } def rgb_to_index(rgb_label): h, w, _ rgb_label.shape index np.zeros((h, w), dtypenp.uint8) for color, class_id in color2id.items(): mask (rgb_label np.array(color)).all(axis-1) index[mask] class_id return index这段代码看起来简单但实际很容易踩坑。第一个坑是 OpenCV 读图默认 BGR 顺序cv2.imread读进来后rgb_label的通道顺序是反的你直接比对 (128, 64, 128) 永远匹配不上。第二个坑是颜色值可能有 JPEG 压缩带来的微小偏移比如 128 变成 127导致掩码全部漏掉。我一般要求标注导出时必须是 PNG 无损格式且不要用 JPEG 保存标注图。如果你发现某些类别的像素数量明显异常少先从这两点排查。3.2 数据增强不要破坏道路结构的空间关系图像分类里常用的随机裁剪、翻转和旋转在分割任务里要更谨慎。比如随机旋转 90 度对道路场景来说可能把“上方是天空、下方是路面”的空间关系打乱模型学到的先验知识被破坏。我常用的增强策略是水平翻转、小幅旋转±5 度以内、随机缩放0.8 到 1.2 倍、随机亮度和对比度调整。这里我直接用albumentations它对图像和掩码同步变换不用手工保证两套变换一致import albumentations as A train_transform A.Compose([ A.Resize(height512, width512), A.HorizontalFlip(p0.5), A.RandomScale(scale_limit(0.8, 1.2), p0.4), A.RandomBrightnessContrast(brightness_limit0.1, contrast_limit0.1, p0.4), ])注意每一项的p参数是应用概率不要全部设为 1。RandomScale会改变图像和掩码的尺寸如果后面接Resize顺序很重要。我习惯把Resize放在最前面保证后续增强在固定尺寸下进行否则RandomScale之后再Resize等于缩放被抵消了增强效果不明显。对于道路分割我不推荐使用粗粒度的 Cutout 或随机遮挡。因为车辆和行人这类目标本身就很依赖完整轮廓中间挖掉一块模型很可能就把那个区域误判成路面。如果确实需要正则化可以试试在特征图上做 Dropout而不是在输入图像上做剪贴。3.3 自定义 Dataset 与数据加载器PyTorch 里写一个自定义 Dataset 是固定套路。我把它拆成几步初始化时读入所有文件路径__getitem__里加载图像和标签做训练增强最后把标签转成torch.long张量。import os import cv2 import torch from torch.utils.data import Dataset class RoadDataset(Dataset): def __init__(self, img_dir, label_dir, transformNone, color_mapNone): self.img_dir img_dir self.label_dir label_dir self.transform transform self.color_map color_map self.img_names sorted(os.listdir(img_dir)) def __len__(self): return len(self.img_names) def __getitem__(self, idx): img_name self.img_names[idx] img_path os.path.join(self.img_dir, img_name) label_path os.path.join(self.label_dir, img_name.replace(.jpg, .png)) image cv2.imread(img_path) image cv2.cvtColor(image, cv2.COLOR_BGR2RGB) label cv2.imread(label_path, cv2.IMREAD_UNCHANGED) if len(label.shape) 3: label self.rgb_to_index(label) if self.transform: transformed self.transform(imageimage, masklabel) image transformed[image] label transformed[mask] image torch.from_numpy(image.transpose(2, 0, 1)).float() / 255.0 label torch.from_numpy(label).long() return image, labelcv2.imread(label_path, cv2.IMREAD_UNCHANGED)这里必须指定UNCHANGED否则 OpenCV 会把单通道 PNG 读成三通道。如果标签是索引图label.shape是二维直接作为 mask 传给 albumentations如果是彩色图需要先转索引。image.transpose(2, 0, 1)是从 HWC 转 CHW标准化通过除以 255.0 实现我一般不在 Dataset 里再做复杂的 Normalize因为这会影响后续的训练稳定性调试。数据加载器里num_workers建议先从 4 开始batch_size视显存定。pin_memoryTrue可以加速 GPU 拷贝但如果你用的是 Windows 系统并且 workers 经常报错可以先关掉试一下。4. 训练与调参从 loss 曲线到 mIoU4.1 优化器与学习率策略U-Net 训练最常见的优化器选择是 Adam 或 SGD。Adam 收敛快一开始 loss 掉得猛但最后 mIoU 可能不如 SGD 磨得高。我的习惯是先用 AdamW 跑 10 个 epoch把网络结构里的坑都排掉比如显存不足、标签错位确认没有问题后再用带 momentum 的 SGD 重训一遍通常能再多涨 2 到 3 个点。学习率设置上ResNet 系主干一般用 1e-4 起步U-Net 我推荐 2e-4 到 5e-4。配合 poly 学习率衰减策略也就是每个 iteration 都降低学习率公式为lr base_lr * (1 - iter / total_iter) ** 0.9。这种方式比固定步长衰减更适合分割任务后期微调更细腻。from torch.optim import AdamW from torch.optim.lr_scheduler import LambdaLR optimizer AdamW(model.parameters(), lr2e-4, weight_decay1e-4) def poly_scheduler(epoch, warmup_epochs2, total_epochs60, base_lr2e-4): if epoch warmup_epochs: return (epoch 1) / warmup_epochs return (1 - (epoch - warmup_epochs) / (total_epochs - warmup_epochs)) ** 0.9 scheduler LambdaLR(optimizer, lr_lambdalambda epoch: poly_scheduler(epoch))weight_decay我设 1e-4太小起不到约束作用太大会导致模型欠拟合。前两个 epoch 做 warmup让学习率从 0.25 倍慢慢升到基准值这一步非常关键。如果一开始就用大学习率BatchNorm 的 moving mean 还没稳定训练 loss 会剧烈波动甚至直接产生 NaN。4.2 评估指标mIoU 的计算方式与细节分割任务最通用的评估指标是 mIoU计算每个类别的预测集合与真实集合的交并比然后取平均。很多人以为用 PyTorch 里现成的函数就行但实际上没这么简单自己写几行反而更可控。def compute_miou(pred_logits, target): pred torch.argmax(pred_logits, dim1) ious [] for cls_id in range(pred_logits.shape[1]): pred_mask (pred cls_id) target_mask (target cls_id) intersection (pred_mask target_mask).sum().float() union (pred_mask | target_mask).sum().float() if union.item() 0: ious.append(torch.tensor(1.0)) else: ious.append(intersection / union) return torch.stack(ious).mean().item()这里有个常见争议某个类别在当前 batch 中完全没有出现它的 union 是 0mIoU 应该怎么算我见过两种做法一种是跳过这个类别只对出现的类别算平均另一种是直接视为 1.0因为真实和预测都是空区域完全重合。在道路分割场景里如果一个 batch 里没有行人我会把它按 1.0 处理否则验证集指标会被这种“偶然缺类”的 batch 拉低显得模型很差误导调参。4.3 训练循环完整骨架训练循环的核心就三步前向算 loss、反向传播、更新权重。但需要注意model.train()和model.eval()的切换以及验证阶段需要关闭梯度计算否则显存会被验证过程占满。train_loader DataLoader(train_set, batch_size8, shuffleTrue, num_workers4) val_loader DataLoader(val_set, batch_size8, shuffleFalse, num_workers4) model UNet(n_classes5).cuda() criterion nn.CrossEntropyLoss(ignore_index-1) optimizer AdamW(model.parameters(), lr2e-4) scheduler LambdaLR(optimizer, lr_lambdalambda epoch: poly_scheduler(epoch)) best_iou 0.0 for epoch in range(60): model.train() train_loss 0.0 for images, labels in train_loader: images, labels images.cuda(), labels.cuda() optimizer.zero_grad() logits model(images) loss criterion(logits, labels) loss.backward() optimizer.step() train_loss loss.item() model.eval() val_iou 0.0 with torch.no_grad(): for images, labels in val_loader: images, labels images.cuda(), labels.cuda() logits model(images) val_iou compute_miou(logits, labels) val_iou / len(val_loader) scheduler.step() if val_iou best_iou: best_iou val_iou torch.save(model.state_dict(), unet_road_best.pt)ignore_index-1是一个很实用的设置。如果你的标签图里有未标注区域把那些像素的值设成 -1交叉熵会自动跳过这些像素不影响梯度更新。我认为这个参数比手动做 mask 简单得多。保存模型时只保存state_dict()不要整个模型都 pickle 保存否则换环境时容易因为类的路径不匹配而加载失败。加载时先实例化模型再load_state_dict这是一个让人少掉不少头发的习惯。5. 避坑清单道路分割最常见的五个问题排查5.1 标签错位loss 小但预测图整体错位现象训练 loss 降得很顺利验证 mIoU 也不差但把预测图叠到原图上一看整辆车或整条路的边缘都朝一个方向偏移了几像素像是标签和图像没对齐。原因数据预处理里图像 resize 和标签 resize 用了不同的插值方式。比如图像用双线性插值标签用最近邻插值缩放尺度相同时还好但只要图像和标签的尺寸不是严格同步就会错位。另一个常见原因是某些标注工具导出时图像和标签的坐标系原点不一致导致裁剪时左右、上下偏差。解决我用的是 albumentations 同步变换基本能避免这类问题。如果是自写代码一定保证图像和标签共享同一个变换参数矩阵比如cv2.warpAffine时同一个 M 矩阵必须同时作用于图像和标签。标签千万不能用双线性插值因为类别编号是离散的插值会生成不存在的类别比如类别 1 和类别 2 之间插出 1.2。5.2 loss 是 NaN或者头几个 epoch 就崩现象训练刚开始loss 显示 nan或者第一轮还没跑完直接报错CUDA error: device-side assert triggered。原因最常见的是标签值超出了模型输出类别数n_classes。比如模型n_classes5但标签图里有像素值是 6交叉熵计算时越界触发 CUDA assert。第二个原因是 batch size 太小加上 BatchNorm 不稳定梯度爆炸导致参数变成 NaN。解决第一步先统计标签的像素取值集合确认最大值小于n_classesnp.unique(label)看一遍。第二步在网络的forward输出后加一个断言assert output.shape[1] target.max().item()只保留到调试阶段。如果是 BatchNorm 问题把 batch size 提到 8 以上或者把 BatchNorm 换成 GroupNorm通常立竿见影。5.3 道路占绝大多数车辆和行人根本学不出来现象训练结束后单看 mIoU 有 85%但打开预测图所有非道路目标要么消失要么糊成一片。原因这是典型的类别不平衡。道路像素占图中 70% 以上交叉熵损失里道路类贡献的梯度远大于车辆和行人网络只要把道路分对loss 就很低小目标成了可以被牺牲的部分。解决一是调整损失权重给稀少类别更高的权重比如torch.class_weight根据像素频率的逆来设置。二是直接把dice_loss的权重调大因为 Dice 对每个类别单独计算。三是检查数据增强里是否加了过多的随机缩放车辆在缩小后只有十几个像素经过连续池化后信息基本丢失。一个小技巧是把训练图像裁成多个 256×256 patch每个 patch 尽量包含一张图里的小目标让模型更频繁地看到行人这类类别。5.4 推理结果出现棋盘格或密密麻麻的小洞现象模型在训练集上效果不错测试视频里道路中间却出现棋盘格状的小块或者单个车辆目标内部有孔洞。原因棋盘格通常来自转置卷积的叠加尤其是当上采样层数多、通道数大时转置卷积会对某些频率产生过强的响应。道路本身的纹理也很复杂模型对高置信度的道路区域内部做 argmax 时可能因为相邻像素特征不一致而输出不同类别。孔洞则说明分类器对目标内部特征的把握不够把车辆中间几块误判成其他类别。解决上采样替换成nn.Upsample(scale_factor2, modebilinear, align_cornersTrue)再接卷积而不是用转置卷积。虽然理论上转置卷积可学习但在小数据集上容易过拟合棋盘格几乎都是它带来的。推理后处理方面可以加一个简单的多数投票对预测图的局部窗口取众数过滤孤立点。用 OpenCV 的modeFilter或morphologyEx都行但注意不要过度处理否则会把车辆和行人边缘磨平。5.5 显存不够batch size 调小后训练效果断崖式下降现象batch size 从 16 降到 4 之后同样的 epoch 数验证 mIoU 掉了十来个点。原因BatchNorm 在小 batch 下统计均值方差抖动厉害模型每个 step 看到的归一化分布都不一样训练不稳定。很多人误以为减 batch size 只是慢一点实际上对分割模型影响很大因为分割模型的 BatchNorm 层很多。解决如果显存只允许 batch size 4建议把所有BatchNorm2d替换成GroupNorm一般分组数设为 32 或 8效果接近大 batch 下的 BatchNorm。另一个做法是使用梯度累积每 4 个小 batch 累积一次梯度再更新等价于 batch size 翻 4 倍但 BN 的统计仍然是基于小 batch 的所以并不能完全解决问题。最彻底的办法是改用预训练编码器比如用 ResNet34 作为 encoder 的 U-Net 变体显存占用小很多而且收敛更快。6. 推理优化与模型落地从单张图到连续帧6.1 滑动窗口推理与大图切块实际落地时你拿到的图片可能不是 512×512而是 1920×1080 的车载截图。直接 resize 到 512×512 会丢失大量细节小目标直接糊掉。我的做法是滑窗推理把原图切成多个 512×512 的块块之间重叠 64 像素每个块独立预测最后把预测结果拼回原图重叠区域取 argmax 或多数投票。这样可以保留原始分辨率代价是推理次数增加。重叠区域的处理我一般直接取第一个块的预测结果不做融合因为融合逻辑复杂且容易在边界引入新问题。如果你追求更平滑的结果可以对重叠部分按像素距离做加权平均但实测差异不大没必要为此增加代码复杂度。6.2 从 PyTorch 到 TorchScript 导出训练完成后如果要把模型放到 Python 服务或 C 端侧可以用 TorchScript 导出省掉 PyTorch 运行时。导出前务必进入model.eval()模式否则 BatchNorm 的 moving average 不会被固定导出的模型结果会不稳定。model UNet(n_classes5) model.load_state_dict(torch.load(unet_road_best.pt)) model.eval() dummy torch.randn(1, 3, 512, 512) traced torch.jit.trace(model, dummy) traced.save(unet_road.pt)导出后用traced(images)验证一次和原模型逐像素对比误差应该在1e-5以下。如果误差很大大概率是torch.jit.trace路径里有条件分支或动态循环U-Net 的静态卷积结构一般不会有这个问题。从那以后我每次训练完都会强制走一遍这个导出流程确保交付出去的模型不是只在训练代码里能跑的“黑匣子”。希望帮到你。本文还有配套的精品资源点击获取
