CNN交通标志分类实战:从GTSRB数据到PyTorch部署
简介围绕智慧交通场景中交通标志识别这一典型任务资源以卷积神经网络CNN为主轴提供一套可直接运行的项目实践方案适合具备基础Python知识、希望系统学习图像分类完整流程的初学者与开发者。压缩包共包含2000个文件整体约201.95MB其中1994张PNG图片构成多类别交通标志样本集3个XML文件用于标注或配置信息3个Python脚本覆盖模型训练和单张图片预测两大环节。资源内给出了可直接调用的训练命令与测试命令用户能够基于train.py在自定义数据集上完成CNN模型训练再通过predict.py对任意指定图片进行实时分类验证从数据准备到结果输出形成清晰闭环。目前已有78人学习下载整套资料样本组织规范、代码量精炼既适合作为课程设计或毕业设计的参考实现也能帮助学习者直观理解CNN在智慧交通场景中的落地方式。1. 为什么交通标志分类先拿 CNN 开刀从一帧误判说起夜间行车记录仪里前方限速 40 的蓝色圆牌被识别成限速 60哪怕只差 0.3 秒判断巡航策略就会出错。智慧交通里的视觉识别最现实、最典型的落地场景不是车道线也不是行人而是交通标志分类。它类别固定、目标尺寸小、环境光照变化大正好把卷积神经网络CNN的三板斧——局部感受野、权值共享、层次化特征提取——全部用上。如果你手头只有一套公开数据集和一台普通 GPU想用人工智能的方式做一个完整项目实践交通标志分类是性价比最高的切入点数据好找、模型不深、训练可控又能真实踩一遍从数据清洗到部署的工程流程。这篇笔记就按我实际做的顺序把数据加载、模型搭建、训练调参和那些坑位讲清楚。2. 数据集与预处理把 GTSRB 的照片变成模型能吃的张量2.1 数据来源与目录结构先弄清楚你手里有什么最常见的交通标志数据集是德国 GTSRB43 个类别包含限速、禁止超车、让行、停止等标志。它的原始目录按Train/0/,Train/1/这样分文件夹文件夹名就是类别 ID每个文件夹里有几百到一千多张不等的图片。还有一个Test/目录放测试集但测试集的标注单独放在一个 CSV 文件里不是图片文件名带标签。我个人做项目时不会直接拿原始 CSV 就开跑而是先把目录结构理成一个统一的清单图片绝对路径 类别 ID。这一步花不了五分钟但能避免后面写 Dataset 加载类时到处补丁。# 生成训练集与测试集的路径-标签清单 # 训练集目录名即标签 find Train -type f -name *.ppm | awk -F/ {print $0, $2} train_list.txt # 测试集从 CSV 里截取第二列文件名和第一列标签 awk -F; NR1{print Test/$2, $1} Test.csv test_list.txt # 看一眼生成的清单 head -5 train_list.txt这段命令是把 PPM 格式图片的路径和类别 ID 拉成两列方便 Python 端直接按行读。注意 GTSRB 的图片是.ppm后缀有些图像库不直接支持后面读取时要转成 RGB。awk -F;是处理 CSV 的常见做法因为 GTSRB 的标注文件用分号分隔不是逗号。如果你是做人工智能课程设计或毕业设计想省事一点也可以直接用torchvision.datasets.ImageFolder读Train/目录它自动按子目录名映射标签。但测试集还是要手动处理。2.2 自定义 Dataset 加载类从文件名到 (image, label) 对直接用ImageFolder的好处是代码短坏处是预处理灵活性差。交通标志识别里几乎必做的一件事是裁剪标志区域因为原始图像里标志只占中间一小块周围有大量背景。GTSRB 官方标注里提供了 ROI 坐标如果不做裁剪模型会把背景也学进去测试时换个场景就翻车。写一个自定义Dataset类把路径读取、裁剪、缩放、灰度转换全收进去。import os import pandas as pd from PIL import Image import torch from torch.utils.data import Dataset from torchvision import transforms class GTSRBDataset(Dataset): 读取 GTSRB 并返回 (tensor_image, label) csv_path: 包含 path, label, roi_x1, roi_y1, roi_x2, roi_y2 列 def __init__(self, csv_path, size(48, 48), use_grayTrue): self.df pd.read_csv(csv_path, sepr\s, headerNone, names[path, label, roi_x1, roi_y1, roi_x2, roi_y2]) self.size size self.use_gray use_gray def __len__(self): return len(self.df) def __getitem__(self, idx): row self.df.iloc[idx] img Image.open(row[path]).convert(RGB) # 按 ROI 裁剪去掉周围街道背景 x1, y1, x2, y2 int(row[roi_x1]), int(row[roi_y1]), int(row[roi_x2]), int(row[roi_y2]) img img.crop((x1, y1, x2, y2)) # 缩放到统一尺寸 img img.resize(self.size, Image.BILINEAR) # 灰度化交通标志本身颜色有区分度但灰度能削弱光照影响 if self.use_gray: img img.convert(L) # 转成张量并归一化到 [0,1] tensor transforms.ToTensor()(img) # 灰度图只有一个通道重复三次便于用预训练权重 if self.use_gray and tensor.shape[0] 1: tensor tensor.repeat(3, 1, 1) label int(row[label]) return tensor, label这里有几个关键设计裁剪用的是官方 ROI 而非全图缩放这能直接让基线准确率提升五到八个百分点灰度转换后我又重复成三通道纯粹是为了以后想加载在 ImageNet 上预训练好的权重时不报通道数错误。标签直接取字符串 int因为 GTSRB 的类别 ID 本身就是 0 到 42 的整数。做项目时我不建议一开始就把数据增强加进 Dataset 里——先跑一个干净版本确定基线再一层层加东西这样出了问题你能判断是哪一步引入的。2.3 预处理策略灰度、归一化与尺寸的取舍很多人拿交通标志直接套 ImageNet 的平均值和方差做归一化但 GTSRB 是德国的高速公路实拍场景像素分布和 ImageNet 差不少。我习惯先算自己数据集的均值和方差虽然代码多两步效果更稳。import numpy as np from tqdm import tqdm def compute_mean_std(dataset, batch_size128): loader torch.utils.data.DataLoader(dataset, batch_sizebatch_size, shuffleFalse) channel_sum np.zeros(3) channel_sq_sum np.zeros(3) total_pixels 0 for images, _ in tqdm(loader): # images shape: [B, 3, H, W] batch_pixels images.shape[0] * images.shape[2] * images.shape[3] channel_sum images.sum(axis(0, 2, 3)).numpy() channel_sq_sum (images ** 2).sum(axis(0, 2, 3)).numpy() total_pixels batch_pixels mean channel_sum / total_pixels std np.sqrt(channel_sq_sum / total_pixels - mean ** 2) return mean, std mean, std compute_mean_std(train_dataset) print(fDataset mean: {mean}, std: {std})参数说明这一步算出来的 mean 和 std 要保存下来测试时用同一组数不能测试集单独算否则分布不一致。尺寸我习惯选48x48够小、训练快同时保留标志内部的文字轮廓。如果你想追求更高精度上64x64或96x96都可以但训练时间会明显增加。预处理策略的取舍根源在于交通标志识别本质是形状识别任务颜色虽有语义红色代表禁令、蓝色代表指示但在复杂光照下不可靠。灰度化可以剔除颜色域的信息量但对暗部细节更敏感所以配合归一化一起用。如果你发现灰度后模型训练不收敛大概率是归一化参数算错了回去检查std里有没有零值。3. 模型设计从 LeNet 基线出发搭一个能跑的 CNN3.1 为什么基线模型选 LeNet-5 而不是 ResNet很多第一次做人工智能项目的同学上来就选 ResNet-50理由是准确率高。但对交通标志分类这个任务ResNet-50 是明显的过度设计GTSRB 图像分辨率小标志结构简单深网络的收益很有限反而带来训练慢、显存占用高、调参难的问题。CNN 基本结构里最经典的 LeNet-5 反而合适作为基线它只有几层卷积和全连接十几分钟就能训练完有了结果再判断是否需要升级模型。当然LeNet-5 也有它的短板没有批量归一化BatchNorm深层网络容易梯度消失感受野偏小对大尺寸标志的边缘信息不够敏感。所以实践里我会在 LeNet 骨架上加两个改进加入 BatchNorm 层以及把激活函数从 ReLU 换成 LeakyReLU防止神经元死亡。从业务角度讲智慧交通场景下交通标志分类更看重的是推理速度和模型体积不是单纯追那几个百分点的精度。一个能在 Jetson Nano 或树莓派上跑到 30 FPS 的轻量 CNN比一个需要高端 GPU 才能跑 ResNet 实际得多。3.2 PyTorch 实现一个交通标志分类 CNN下面是我实际用的模型结构基于 LeNet 的骨架改写输入48x48的 RGB 图像输出 43 类概率。import torch.nn as nn class TrafficSignCNN(nn.Module): def __init__(self, num_classes43): super().__init__() # 第一段卷积激活池化 self.features nn.Sequential( nn.Conv2d(3, 32, kernel_size3, padding1), # 48x48 - 48x48 nn.BatchNorm2d(32), nn.LeakyReLU(0.1, inplaceTrue), nn.MaxPool2d(2), # 48x48 - 24x24 nn.Conv2d(32, 64, kernel_size3, padding1), # 24x24 - 24x24 nn.BatchNorm2d(64), nn.LeakyReLU(0.1, inplaceTrue), nn.MaxPool2d(2), # 24x24 - 12x12 nn.Conv2d(64, 128, kernel_size3, padding1), # 12x12 - 12x12 nn.BatchNorm2d(128), nn.LeakyReLU(0.1, inplaceTrue), nn.MaxPool2d(2), # 12x12 - 6x6 ) # 分类头拉平后过两层全连接 self.classifier nn.Sequential( nn.Flatten(), nn.Linear(128 * 6 * 6, 512), nn.ReLU(inplaceTrue), nn.Dropout(0.3), # 防止过拟合 nn.Linear(512, num_classes), ) def forward(self, x): x self.features(x) x self.classifier(x) return x逻辑说明这里没有直接照搬 LeNet 的5x5卷积换成了3x3加 padding保证特征图尺寸不缩水信息保持更好。三个卷积段的通道数从 32 递增到 128每段后面都接 BatchNorm 和池化形成“特征提取 降采样”的模式。全连接层的 Dropout 设 0.3因为交通标志数据集不大全连接层容易过拟合。参数说明kernel_size3, padding1是卷积标配保持分辨率不降MaxPool2d(2)把尺寸减半扩大感受野的同时减少计算量Dropout(0.3)只在全连接层加卷积层不加保留空间特征完整性。输出层直接用nn.Linear(512, num_classes)不加激活因为后面会用交叉熵损失它内部自带 Softmax。我一般会用torchsummary或直接打印模型结构检查一下各层的输出维度写错维度是最常见的低级错误。3.3 损失函数与输出层设计43 类分类的细节分类任务首选损失函数是交叉熵PyTorch 里的nn.CrossEntropyLoss把 Softmax 和损失计算打包好了。但交通标志分类有个特殊点这个类别分布极其不均衡。GTSRB 里类别 0限速 20只有一百多张类别 1限速 30有上千张。如果用标准交叉熵模型会偏向样本多的类别少数类准确率掉得很惨。常见做法是给损失函数加类别权重权重值反比于样本数。GTSRB 每类的样本数在 CSV 里没有现成的需要自己统计一下。统计后可以直接用torch.tensor构造权重传给损失函数。import torch from collections import Counter # 统计训练集中每个类别的样本数 label_counter Counter(train_dataset.df[label]) num_classes 43 # 权重 总样本数 / (类别数 * 该类样本数)平滑一下防止极端值 total sum(label_counter.values()) class_weights [total / (num_classes * label_counter[i]) for i in range(num_classes)] class_weights torch.tensor(class_weights, dtypetorch.float32) criterion nn.CrossEntropyLoss(weightclass_weights)参数说明CrossEntropyLoss的weight参数要求是一个长度等于类别数的张量它会在每个样本的损失上乘以对应类别权重。这里的计算方式保证了各类权重之和为 1近似少数类权重高、多数类权重低模型不会一边倒。如果样本极度不均衡还可以配合Focal Loss使用但 GTSRB 用class_weights已经足够。输出层不加 sigmoid 或 softmax 的原因是CrossEntropyLoss内部做了严格的数值稳定处理直接输出 logits未经过激活的原始分数给损失函数比手动加 Softmax 再算损失在数值上更稳反向传播也更快。4. 训练与调参让准确率从 80% 涨到 95% 的关键操作4.1 优化器选择与学习率调度交通标志分类的训练里能明显拉开差距的往往不是模型结构而是优化器和学习率设置。我刚跑这个数据集时用 Adam默认学习率 0.001到第 30 轮卡在 94% 上不去。后来换成 SGD Momentum 搭配余弦退火学习率同样轮数次到了 97%。原因不玄学Adam 的自适应学习率让后期参数更新步长变小难以精细收敛SGD 的全局学习率配合动量项能更彻底地搜索损失曲面的低谷。optimizer torch.optim.SGD(model.parameters(), lr0.01, momentum0.9, weight_decay5e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max30, eta_min1e-4)参数说明lr0.01比 Adam 默认的 0.001 高一倍因为 SGD 没有自适应修正初始步长太小会训练很慢momentum0.9是标准动量设置减少震荡weight_decay5e-4是 L2 正则防止全连接层记住训练集。CosineAnnealingLR的T_max30表示 30 个 epoch 内从初始学习率余弦降到eta_min一轮完整训练下来学习率呈平滑曲线前段快学后段细调。如果你不想手动调学习率也可以用ReduceLROnPlateau它会监视验证集损失下降多少连续几个 epoch 没下降就自动把学习率减半。但我的实践经验是余弦退火的最终效果更稳定且不需要频繁干预。训练时每轮结束都记录训练集和验证集的准确率。如果训练集 99% 而验证集只有 92%这是过拟合信号回去改 Dropout 或数据增强如果两者都低先怀疑学习率是不是太高或太低看一眼前 5 轮的 loss 曲线再判断。4.2 数据增强用小手段治过拟合交通标志分类的过拟合主要来自背景干扰和位置偏移。原始图片经过 ROI 裁剪后标志基本居中但实际部署时相机角度有偏差、标志可能被部分遮挡。数据增强是让模型对这些变化不敏感的最直接手段。from torchvision import transforms train_transforms transforms.Compose([ transforms.RandomAffine(degrees10, translate(0.1, 0.1), scale(0.9, 1.1)), transforms.ColorJitter(brightness0.2, contrast0.2), transforms.ToTensor(), transforms.Normalize(meanmean, stdstd) ])参数说明RandomAffine的degrees10只让图案旋转 ±10 度因为交通标志的安装方向基本固定转太大角度反而脱离真实分布translate(0.1, 0.1)在 10% 范围内随机平移模拟标志不在画面正中心的情形scale(0.9, 1.1)缩放 10% 模拟距离变化。ColorJitter的亮度对比度各 0.2模拟不同天气和日照强度。注意这里和 2.3 节里的基础预处理不冲突先把裁剪、缩放、灰度化放在 Dataset 里做数据增强再叠加在已经规整化的图像上。顺序别反否则增强时会把 ROI 区域外的背景噪声一起放大。使用数据增强后训练集 loss 会比之前高一点这是预期内的——模型在学习更泛化的特征而非背诵特定图片。重点看验证集 acc 是否提升如果验证集也跟着掉说明增强强度太大把degrees降到 5 或translate降到 0.05 再试。4.3 训练循环与验证指标不仅要看 acc 还要看每类召回一个完整的训练循环要同时处理好训练和验证两种模式。PyTorch 里用model.train()和model.eval()切换这两行写错会导致 BatchNorm 和 Dropout 在验证时“偷看”训练分布或者验证时还随机丢弃神经元结果完全不稳定。我一般会单独用一个函数封装训练循环每个 epoch 结束时跑一次验证集并打印 top-1 准确率。但只看整体准确率远远不够还得看混淆矩阵里哪些类容易混淆。def train_one_epoch(model, dataloader, optimizer, criterion, device): model.train() running_loss 0.0 correct 0 total 0 for images, labels in dataloader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) # 前向传播 loss criterion(outputs, labels) # 计算损失 loss.backward() # 反向传播 optimizer.step() # 更新参数 running_loss loss.item() * images.size(0) _, pred torch.max(outputs, 1) total labels.size(0) correct (pred labels).sum().item() epoch_loss running_loss / total epoch_acc correct / total return epoch_loss, epoch_acc逻辑说明前向传播得到outputs后直接和labels算损失outputs是 [batch, 43] 的原始分数矩阵不需要手动argmax后再算准确率。准确率是拿pred每张图的最大分数索引与labels逐像素比较后求和除以总数。用total维护当前累计样本量是为了算加权平均损失避免最后一个 batch 数据量不同导致平均失真。验证集除了整体 acc一定要按类别统计召回率。GTSRB 中“限速 30”“限速 50”“限速 80”这几个类别细节相似召回率经常互相拖累。如果某一类召回率明显低于其他类别优先检查该类的训练样本数量够不够再考虑是否动用数据增强里的特定变换。验证集如果不单独切可以直接用 GTSRB 官方的Test/目录。我没有用它做训练时验证而是最后提交前测一次避免验证集信息泄漏到模型选择里。5. 避坑清单交通标志分类常见的 5 个翻车现场5.1 现象测试集准确率比验证集低 10 个百分点我在跑 GTSRB 时踩过最深的坑验证集 96%提交官方测试集只有 86%。原因出在预处理不一致——训练和验证用了我上面写的自定义 Dataset里面的裁剪、灰度、归一化参数和测试集加载代码不完全一样。测试集没有经过相同规律的 ROI 裁剪而是用了整图缩放。这个坑极隐蔽因为训练时没报什么错误loss 曲线也正常。解决办法是给三个数据集训练、验证、测试写同一个预处理函数不要复制粘贴代码。我把 2.2 节的__getitem__抽出来做成一个独立的preprocess_image()函数三个 Dataset 都调它从根上避免不一致。5.2 现象训练 loss 不掉准确率一直在 2% 左右随机水平原因通常是标签错位。GTSRB 的类别 ID 是 0 到 42但如果你用了ImageFolder读取它的类别排序是按文件夹名的字母序排的不是按数字序号。文件夹Train/0/会被排在Train/10/后面导致 label 对应错位。解决打印出dataset.class_to_idx看一眼映射关系或者干脆不用ImageFolder直接用我上面写的自定义 Dataset 读取 CSV 清单。CSV 里label列是官方标注不会有排序歧义。5.3 现象验证集准确率上去了但某些类别完全不预测具体表现是混淆矩阵里某一行全是 0。原因是不均衡。GTSRB 中有些类别只有几十张训练图模型学到了平均值上就直接忽略它们。我一开始没加 class weight所有类别权重都是 1结果“限速 20”“限速 120”这类少样本类别一个也没预测对。解决按 3.3 节的方式先统计类别样本数给CrossEntropyLoss传入weight。加了权重后少样本类别的召回率从 0% 涨到 60% 以上整体准确率反而提升 1 到 2 个百分点因为模型不再盲目偏向多数类。5.4 现象训练一段时间后 loss 突然变成 NaN原因大概率是梯度爆炸。初始学习率过大、BatchNorm 的 momentum 设置不当、或者数据里有异常像素值都可能触发。我在跑灰度图时遇到过因为局部像素值和高斯初始化权重相乘后某些通道的梯度突破数值上界。解决方式分两步先把学习率从 0.01 降到 0.005 看是否复现再用torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm5)做梯度裁剪。梯度裁剪是防御武器正常训练时不会触发但能让 NaN 彻底消失。5.5 现象模型能分类但部署到嵌入式设备上慢得没法用我从 PyTorch 训练完直接导出.pt文件放到 Jeston Nano 上推理一张图要 200 毫秒完全达不到实时。原因是模型结构和推理框架没做优化还停留在训练模式的代码路径。解决方式是模型量化。PyTorch 的torch.quantization可以把 float32 权重压到 int8体积缩为原来的四分之一推理速度提升三到五倍精度损失在1%以内。量化后再用torch.jit.script或者ONNX导出部署时就不需要依赖完整的 PyTorch 环境。有条件的话再配合 TensorRT速度还能再翻一倍。6. 从 95% 到能用混淆矩阵、可解释性与边缘部署模型在验证集上达到 95% 以上只能算是“毛坯房”离真正能放进智慧交通系统还有一段路。这一步我会做三件事画混淆矩阵找到系统性的混淆模式用可解释性工具确认模型到底在看什么最后把模型裁剪成能在边缘设备实时跑的版本。混淆矩阵是诊断工具里性价比最高的。它比准确率直观得多——直接暴露哪两类图片容易搞混。交通标志分类中最经典的混淆是对“限速 30”和“限速 50”的误判两者的圆形轮廓几乎一样区别只在数字。如果混淆集中在这种细节差异简单加卷积层深度未必有用更好的方向是用数据增强里的RandomAffine加大旋转角度让模型对数字形状更敏感。在 PyTorch 里画混淆矩阵可以用sklearn.metrics.confusion_matrix配合seaborn可视化。可解释性上我常用的是 Grad-CAM它能生成一张热力图标注模型分类时关注图像中的哪些区域。对第 2 章里的灰度图做 Grad-CAM 特别值得看如果热力区不在标志区域而是散布在背景上说明模型学了背景特征而不是标志本身这种模型在真实场景里往往一换环境就崩。PyTorch 有现成的pytorch_grad_cam库不用自己实现反向传播钩子。这里顺带提一句“离线解释器”把热力图存成本地图片方便在报告里展示模型的注意点比纯文字说明有说服力。部署这一步我会先把模型从 PyTorch 转成 ONNX 格式再根据目标设备决定是直接用 ONNX Runtime 还是转成 TensorRT。转 ONNX 时有几个注意点输入输出张量的动态维度要标清楚模型里不能用 Python 原生控制流比如if batch_size 1这种PyTorch 的 JIT 能编译的模型才适合导出。import torch.onnx model.eval() dummy_input torch.randn(1, 3, 48, 48) torch.onnx.export( model, dummy_input, traffic_sign_cnn.onnx, export_paramsTrue, opset_version11, input_names[input], output_names[output], dynamic_axes{input: {0: batch_size}, output: {0: batch_size}} )这块的经验是先锁定边缘设备的推理框架再选导出路径。ONNX Runtime 适合第一步快速验证TensorRT 适合真的把速度压到底。最后说一个我的习惯每做完一个阶段就把混淆矩阵和 Grad-CAM 结果截图存档。因为即使你转了 ONNX、做了量化中间只要改了预处理策略这些存档能帮你快速定位问题是出在模型本身还是前后处理。做完这套流程一个能扛住实际路测的交通标志分类器才算真正落地。希望帮到你。本文还有配套的精品资源点击获取