简介:基于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×3,padding 为 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_size=3, padding=1), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True), nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True) ) 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,否则训练曲线会像心电图一样上下跳。inplace=True是为了省显存,反向传播时 ReLU 不需要保留输入,能省一点是一点。
编码器部分就是串行调用这些块和池化层。我用 MaxPool2d 而不是 Conv2d stride=2 做下采样,原因是池化没有可学习参数,训练更稳定,且不容易丢失低层位置信息。
2.2 解码器:上采样与跳跃连接
编码器把一张 512×512 的图逐步压缩到 32×32 的高维特征,但分割要求输出和原图一样大小的逐像素类别,所以需要一个解码器把空间分辨率逐层恢复。解码器的核心操作是转置卷积上采样,再把对应编码器层的特征图拼回来。
这个“拼回来”的操作就是 U-Net 的跳跃连接。它解决的是一个很实际的问题:经过四次池化后,高层特征虽然语义信息丰富,但边界和细节已经被抹掉了。跳跃连接直接把浅层的细节特征传给解码器,让网络在恢复空间尺寸的同时,还有足够的信息去修正边缘。
class UNet(nn.Module): def __init__(self, n_classes=8): 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_size=2, stride=2) self.dec1 = DoubleConv(1024, 512) self.up2 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2) self.dec2 = DoubleConv(512, 256) self.up3 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2) self.dec3 = DoubleConv(256, 128) self.up4 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2) self.dec4 = DoubleConv(128, 64) self.out_conv = nn.Conv2d(64, n_classes, kernel_size=1) 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], dim=1)) d2 = self.dec2(torch.cat([self.up2(d1), e3], dim=1)) d3 = self.dec3(torch.cat([self.up3(d2), e2], dim=1)) d4 = self.dec4(torch.cat([self.up4(d3), e1], dim=1)) return self.out_conv(d4)torch.cat的dim=1是在通道维上拼接。注意上采样后的特征图尺寸必须和对应编码器层完全一致,否则拼接会直接报错。如果输入尺寸不是 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.CrossEntropyLoss,PyTorch 里它对logits直接计算,不需要手动做 softmax。对于道路目标这类场景,类别之间数量极不平衡:道路可能占图像面积的 60%,行人可能只占 0.5%。直接用交叉熵,网络会倾向于把所有像素都预测成道路,mIoU 看起来还行,行人这一类直接废掉。
所以我常把交叉熵和 Dice Loss 结合起来。Dice Loss 对前景小目标更友好,计算的是预测区域和真实区域的像素重叠程度。简单实现如下:
def dice_loss(pred, target, smooth=1.0): pred = torch.softmax(pred, dim=1) target_onehot = torch.nn.functional.one_hot(target, num_classes=pred.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), dtype=np.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(height=512, width=512), A.HorizontalFlip(p=0.5), A.RandomScale(scale_limit=(0.8, 1.2), p=0.4), A.RandomBrightnessContrast(brightness_limit=0.1, contrast_limit=0.1, p=0.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, transform=None, color_map=None): 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(image=image, mask=label) 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_memory=True可以加速 GPU 拷贝,但如果你用的是 Windows 系统并且 workers 经常报错,可以先关掉试一下。
4. 训练与调参:从 loss 曲线到 mIoU
4.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(), lr=2e-4, weight_decay=1e-4) def poly_scheduler(epoch, warmup_epochs=2, total_epochs=60, base_lr=2e-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_lambda=lambda 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, dim=1) 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 是 0,mIoU 应该怎么算?我见过两种做法,一种是跳过这个类别,只对出现的类别算平均;另一种是直接视为 1.0,因为真实和预测都是空,区域完全重合。在道路分割场景里,如果一个 batch 里没有行人,我会把它按 1.0 处理,否则验证集指标会被这种“偶然缺类”的 batch 拉低,显得模型很差,误导调参。
4.3 训练循环完整骨架
训练循环的核心就三步:前向算 loss、反向传播、更新权重。但需要注意model.train()和model.eval()的切换,以及验证阶段需要关闭梯度计算,否则显存会被验证过程占满。
train_loader = DataLoader(train_set, batch_size=8, shuffle=True, num_workers=4) val_loader = DataLoader(val_set, batch_size=8, shuffle=False, num_workers=4) model = UNet(n_classes=5).cuda() criterion = nn.CrossEntropyLoss(ignore_index=-1) optimizer = AdamW(model.parameters(), lr=2e-4) scheduler = LambdaLR(optimizer, lr_lambda=lambda 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_classes=5,但标签图里有像素值是 6,交叉熵计算时越界触发 CUDA assert。第二个原因是 batch size 太小加上 BatchNorm 不稳定,梯度爆炸导致参数变成 NaN。
解决:第一步先统计标签的像素取值集合,确认最大值小于n_classes:np.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_factor=2, mode='bilinear', align_corners=True)再接卷积,而不是用转置卷积。虽然理论上转置卷积可学习,但在小数据集上容易过拟合,棋盘格几乎都是它带来的。推理后处理方面,可以加一个简单的多数投票:对预测图的局部窗口取众数,过滤孤立点。用 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_classes=5) 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 的静态卷积结构一般不会有这个问题。从那以后我每次训练完,都会强制走一遍这个导出流程,确保交付出去的模型不是只在训练代码里能跑的“黑匣子”。希望帮到你。
本文还有配套的精品资源,点击获取