简介本资源面向自动驾驶感知方向的学习者与算法工程师提供基于Unet-MobileNet的道路信息与车道线图像分割实战代码适合具备一定深度学习基础、希望快速复现语义分割项目的人群。压缩包共约2000个文件以1987张png图像含原图与对应mask标注为主体另含8个py脚本、3个txt说明、1个json配置与1个md文档整体约280.8MB可直接用于训练与验证。代码在常规分割流程上补充了loss、iou、dice、recall、precision及对应mean指标并自动生成各类别与mean的训练、验证曲线图便于观察收敛过程。项目支持一键运行参考指标中mean precision约0.9718、mean recall约0.9300、mean dice约0.9497、mean iou约0.9070效果直观可对照。若需训练自有数据集按readme说明调整即可。目前已有110人学习适合作为分割入门与工程落地的参考方案。1. 自动驾驶道路分割这套 Unet-MobileNet 代码我拆完发现最值钱的是评估指标自动驾驶感知链路里道路信息和车道线分割是最底层也最容易被低估的一环。你可能见过太多「Unet 图像分割」的 demo跑通一张图就敢叫实战但真拿去喂给下游的端到端模型或者规控模块才发现连最基本的 IoU、Dice 都没算清楚更别提按类别分开看 precision 和 recall。这套基于 Unet-MobileNet 的自动驾驶道路信息、车道线图像分割代码我拆下来的第一感受是它把「能跑」和「能评估」这两件事同时做完了。数据集、mask 标注、训练脚本、评估指标、曲线图生成一条龙README 里还写了怎么换自己的数据集。适合谁做自动驾驶感知入门、想拿分割练手、或者需要一个干净 baseline 去改 Unet 结构的人。不适合谁指望开箱即用直接上车规级部署的这套代码是研究向的不是工程交付向的。2. Unet-MobileNet 的选型逻辑为什么不是纯 Unet也不是 DeepLab2.1 编码器换 MobileNet 的真实动机纯 Unet 的编码器是 VGG 风格的堆叠卷积参数量大、推理慢在自动驾驶这种对实时性有要求的场景里拿它做 baseline 可以但拿它做迭代起点就有点笨重。MobileNet 的核心是深度可分离卷积把标准卷积拆成 depthwise 和 pointwise 两步参数量和计算量都压下来一个量级。这套代码用 MobileNet 做编码器Unet 的解码器结构保留本质上是一个轻量化的语义分割网络。我一般会跟人说这个组合的甜点区在「你有一块不算顶级的显卡但想跑一个能看的分割模型」。MobileNet 作为 backbone 还有一个隐性好处预训练权重好找ImageNet 上训过的特征提取器拿来初始化收敛比从 scratch 训快得多。热搜里有人搜「scratch自动驾驶」如果你真从零训这套代码也能跑但 epoch 得往上加学习率策略也得调后面避坑章节会细说。2.2 Unet 解码器的跳跃连接在这套代码里怎么落地Unet 之所以在分割任务里经久不衰核心就是编码器和解码器之间的 skip connection。编码器逐层下采样语义信息越来越强但空间分辨率越来越低解码器逐层上采样恢复分辨率但光靠上采样会丢细节。skip connection 把编码器对应层的特征直接拼到解码器上让浅层的边缘、纹理信息和深层的语义信息融合。这套代码里MobileNet 的某些层输出被抽出来作为 skip 的输入解码器对应层做 concat。具体抽哪几层、通道数怎么对齐是改网络结构时最容易翻车的地方。我见过有人直接把 MobileNet 所有层都拿来 skip结果显存爆了通道数对不上训练直接报维度错误。常见做法是选 4 到 5 个关键 stage 的输出通道数用 1x1 卷积压到统一维度再拼。2.3 道路信息和车道线两类目标的类别设计从数据集里的 mask 文件命名看这套数据是二分类或者多分类的分割任务。道路信息和车道线在视觉特征上差异明显道路是大面积连续区域车道线是细长结构。这两类目标放在一个模型里训损失函数会面临类别不平衡的问题——道路像素远多于车道线像素。代码里评估指标按类别分开算 precision、recall、IoU、Dice这个设计很关键因为 mean 指标会掩盖小类别的糟糕表现。如果你只看 mean IoU 0.9070 觉得挺好但车道线那一类的 IoU 只有 0.8334说明细长结构的分割还有提升空间。3. 一键跑通训练从 dataset.json 到指标曲线图3.1 环境依赖和目录结构确认拿到代码包第一件事不是急着跑是先看目录结构。dataset.json 是数据索引文件里面记录了每张原图和对应 mask 的路径映射。README.md 是操作手册换自己数据集的关键步骤都在里面。那些*_mask.png是标注文件命名规则是原图名加_mask后缀。环境依赖常见的是 PyTorch、torchvision、numpy、matplotlib、opencv-python、tqdm。我一般会先建一个干净的 conda 环境避免和系统里的包打架。conda create -n unet_mobilenet python3.8 -y conda activate unet_mobilenet pip install torch torchvision numpy matplotlib opencv-python tqdm这里 Python 版本选 3.8 是因为很多分割库对 3.9 的支持还不稳定3.8 是踩坑最少的选择。torch 版本根据你的 CUDA 版本去官网查对应命令别直接pip install torch装成 CPU 版训练速度会让你怀疑人生。3.2 dataset.json 的结构和自定义数据集替换dataset.json 本质是一个列表每个元素包含原图路径和 mask 路径。换自己数据集时你需要按同样的格式生成这个文件。常见做法是写一个脚本遍历你的图片目录按命名规则配对。import json import os import glob # 原图和 mask 存放的根目录 img_dir your_data/images mask_dir your_data/masks # 假设原图是 .jpgmask 是 _mask.png img_files sorted(glob.glob(os.path.join(img_dir, *.jpg))) dataset [] for img_path in img_files: base os.path.basename(img_path).replace(.jpg, ) mask_path os.path.join(mask_dir, base _mask.png) if os.path.exists(mask_path): dataset.append({ image: img_path, mask: mask_path }) else: print(f警告{img_path} 没有对应 mask已跳过) with open(dataset.json, w) as f: json.dump(dataset, f, indent2) print(f共生成 {len(dataset)} 条数据索引)这段脚本的逻辑很直白遍历原图按命名规则找 mask找到就配对找不到就跳过并打印警告。参数上你只需要改img_dir和mask_dir两个路径。注意 mask 的像素值必须和类别定义一致比如道路是 1、车道线是 2、背景是 0如果你的标注工具输出的是 255得先做一次映射否则训练时类别数对不上loss 会直接 NaN。3.3 训练脚本的关键参数和启动命令训练脚本里几个参数决定成败batch size、学习率、epoch 数、输入分辨率。这套代码 README 里给的参考结果是 epoch 100 跑出来的但没说 batch size 和 lr。我一般会从 batch size 4 或 8 起步学习率 1e-4 配 Adam如果显存够就往上加 batch size学习率同步放大。python train.py \ --dataset dataset.json \ --epochs 100 \ --batch_size 8 \ --lr 1e-4 \ --input_size 512 512 \ --num_classes 3 \ --save_dir checkpoints--input_size设成 512x512 是分割任务的常见选择再大显存吃不消再小车道线这种细结构会糊掉。--num_classes根据你的标注来背景加道路加车道线就是 3。--save_dir是权重保存路径建议每 10 个 epoch 存一次防止训练中途崩了没有后悔药。3.4 评估指标的计算逻辑和曲线图生成这套代码的评估部分是我觉得最值得看的地方。它算了 loss、IoU、Dice、Recall、Precision而且每个类别单独算最后再算 mean。IoU 是交集除以并集Dice 是 2 倍交集除以两者之和这两个指标在分割里最常用。Precision 和 Recall 则是从像素分类的角度看查准率和查全率。import numpy as np def compute_metrics(pred, target, num_classes): pred: 模型输出经过 argmax 后的预测 maskshape (H, W) target: 真实 maskshape (H, W) num_classes: 类别数 iou_list, dice_list [], [] precision_list, recall_list [], [] for cls in range(num_classes): pred_cls (pred cls) target_cls (target cls) intersection np.logical_and(pred_cls, target_cls).sum() union np.logical_or(pred_cls, target_cls).sum() pred_sum pred_cls.sum() target_sum target_cls.sum() # IoU交集除以并集并集为 0 时跳过 iou intersection / union if union 0 else float(nan) # Dice2 倍交集除以预测和真实之和 dice 2 * intersection / (pred_sum target_sum) if (pred_sum target_sum) 0 else float(nan) # Precision预测为该类的像素里有多少是对的 precision intersection / pred_sum if pred_sum 0 else float(nan) # Recall真实为该类的像素里有多少被找出来了 recall intersection / target_sum if target_sum 0 else float(nan) iou_list.append(iou) dice_list.append(dice) precision_list.append(precision) recall_list.append(recall) return iou_list, dice_list, precision_list, recall_list这段代码里每个指标都做了除零保护返回 nan 而不是直接崩。实际训练时曲线图就是把这些指标按 epoch 画出来训练集和验证集各一条线。看曲线图的技巧是如果训练集指标一直涨但验证集指标平了或者掉了就是过拟合该加数据增强或者加 dropout 了。如果两条线都平在低位那是欠拟合得加模型容量或者调学习率。4. 避坑与排查这套代码我踩过的五个坑4.1 现象训练 loss 直接 NaN第一个 epoch 就崩原因通常是 mask 像素值超出了 num_classes 范围。比如你标注时背景是 0、道路是 1、车道线是 2但标注工具导出时把车道线存成了 255。模型输出通道数是 3交叉熵损失遇到 255 这个类别索引直接越界。解决在数据加载的 transform 里加一步像素值映射把所有非 0/1/2 的值重新映射到合法范围。或者用np.unique先检查 mask 里到底有哪些像素值确认后再训练。4.2 现象验证集 IoU 比训练集还高看着很美好原因一般是验证集太小或者分布和训练集不一致。如果验证集只有几十张图而且恰好都是简单场景指标虚高很正常。另一个可能是数据增强只加在训练集上验证集用的是原图模型在干净数据上表现更好。解决验证集至少占总数据的 15% 到 20%而且要覆盖不同光照、不同道路类型。数据增强在验证集上可以只做 resize不做随机裁剪和翻转保证评估的一致性。4.3 现象车道线那一类的 recall 只有 0.86但 precision 有 0.96这是典型的漏检多于误检。车道线像素少模型倾向于保守预测宁可少画也不画错。原因通常是损失函数没有对类别不平衡做处理背景像素主导了梯度。解决换带类别权重的交叉熵损失或者用 Dice loss 和交叉熵组合。Dice loss 对类别不平衡更鲁棒因为它直接优化重叠区域。代码里如果只用了交叉熵可以手动加一个 weight 参数给车道线类别更高的权重。4.4 现象换了自定义数据集后训练速度突然变慢很多原因可能是图片分辨率不一致。原始数据集可能都是 512x512你自己的数据有 1920x1080 的dataloader 每次都要 resizeCPU 成了瓶颈。另一个可能是 dataset.json 里路径写错了dataloader 反复重试读文件。解决提前把所有图片统一 resize 到目标分辨率再存一份训练时直接读。检查 dataset.json 里的路径是不是绝对路径或者相对于运行目录的正确相对路径。4.5 现象曲线图里 mean IoU 到 80 epoch 就不涨了但 loss 还在降这是过拟合的典型信号。loss 降是因为模型在记忆训练集的噪声但泛化能力已经到顶。mean IoU 平了说明验证集上的表现不再提升。解决加早停策略验证集 IoU 连续 10 个 epoch 不提升就停。同时加数据增强随机旋转、亮度调整、高斯噪声都是分割任务里常用的。如果数据量实在少考虑用预训练权重冻结编码器先训解码器再解冻全网络微调。5. 进阶技巧用类别权重和 TTA 把车道线 IoU 再往上推一截这套代码跑出来的参考指标里道路那一类的 IoU 是 0.9805车道线是 0.8334差距很明显。如果你想让车道线分割更准有两个方向可以试。第一个是类别权重。在损失函数里给车道线类别更高的权重让模型更关注这个类。具体做法是在交叉熵里传 weight 参数车道线的权重设成道路的 3 到 5 倍。我一般会先跑一版看混淆矩阵确认车道线的漏检比例再决定权重给多大。import torch import torch.nn as nn # 假设类别顺序是0 背景1 道路2 车道线 # 车道线权重给 5道路给 2背景给 1 class_weights torch.tensor([1.0, 2.0, 5.0]).cuda() criterion nn.CrossEntropyLoss(weightclass_weights)第二个是测试时增强TTA。推理时把图片做水平翻转、多尺度缩放分别预测后再把结果平均或者投票。这个技巧不改变训练过程只在推理阶段加计算量通常能涨 1 到 2 个点的 IoU。代价是推理速度变慢适合离线评估或者对实时性要求不高的场景。def tta_predict(model, image): image: 单张图片 tensorshape (1, C, H, W) 返回融合后的预测 mask model.eval() preds [] with torch.no_grad(): # 原始预测 preds.append(model(image)) # 水平翻转预测 pred_flip model(torch.flip(image, dims[3])) preds.append(torch.flip(pred_flip, dims[3])) # 多尺度缩小到 0.75 倍再预测再上采样回来 h, w image.shape[2], image.shape[3] small torch.nn.functional.interpolate(image, size(int(h*0.75), int(w*0.75)), modebilinear) pred_small model(small) preds.append(torch.nn.functional.interpolate(pred_small, size(h, w), modebilinear)) # 平均后取 argmax avg_pred torch.mean(torch.stack(preds), dim0) return torch.argmax(avg_pred, dim1)这段 TTA 代码做了三种变换原始、水平翻转、0.75 倍缩放。每种变换的预测结果都还原到原图尺寸后平均最后取 argmax 得到最终 mask。注意翻转后的预测要再翻回来才能和原始预测对齐这个细节容易漏。多尺度那一路的上采样用 bilinear 就行分割任务里对插值方式不敏感。从那以后我每次换数据集或者改网络结构都强制先跑 5 个 epoch 看 loss 和指标曲线确认没有 NaN 和维度错误再开长训练。这套 Unet-MobileNet 代码的评估部分做得比大多数 demo 都完整按类别看指标这个习惯也是被它逼出来的。希望帮到你。本文还有配套的精品资源点击获取
