GTSRB交通标志识别实战:TSR-master源码解析与CNN训练全流程
简介这份资源是一套基于Python与卷积神经网络实现的交通标志识别项目采用GTSRB数据集面向计算机、人工智能、通信工程、自动化等专业的在校学生与教师也适合作为毕业设计、课程设计或项目立项的参考案例。压缩包共9个文件约311KB包含5个py源码文件、2个csv数据文件、1个xml配置文件和1个md说明文档分别承担模型构建、训练评估、数据读取与预处理等职责结构清晰便于按模块学习。项目围绕交通标志图像分类任务涵盖数据预处理、CNN模型搭建、训练与评估等完整流程代码均经过实际运行验证可直接复现。目前已有224人学习下载适合希望快速上手深度学习图像分类、理解CNN在真实数据集上应用的学习者参考借鉴。1. 从一份能跑通的 GTSRB 源码包说起它到底解决了什么问题交通标志识别这个方向很多同学第一次接触时都会卡在同一个地方数据集下下来了模型结构也看懂了但真要把 GTSRB 里那几万张 32×32 的彩色小图喂进网络、跑出一份能写进答辩 PPT 的准确率曲线中间隔着一堆琐碎的工程活。这份 TSR-master 源码包就是冲着这个断层来的——它把数据预处理、CSV 索引构建、CNN 模型定义、训练循环和评估脚本拆成了独立文件配合 GTSRB 数据集可以直接跑起来。适合的人群很明确做毕设的本科生、上深度学习课要交课程设计的研究生以及刚学完 CNN 卷积神经网络理论、想找一个完整项目练手的人。它不教你反向传播怎么推导但它给你一条从原始图片到分类结果的完整链路这条链路本身就是最值钱的部分。2. 拆开 TSR-master文件职责与数据流的对应关系2.1 每个 .py 文件在流水线里干什么拿到一个源码包我习惯先不急着跑而是把文件列表和 README 对一遍搞清楚数据从哪进、从哪出。TSR-master 的结构不算复杂但每个文件的职责边界很清晰理解了这个边界后面改代码才知道该动哪里。文件职责输入输出Preprocessing.py图片预处理与数据集划分GTSRB 原始图片目录train_data.csv / test_data.csvTSRInput.py数据读取与批次生成CSV 索引文件 图片路径批量张量TSRCnn.pyCNN 网络结构定义输入张量分类 logitsTSRTrain.py训练主循环模型 训练数据权重文件TSREval.py模型评估权重文件 测试数据准确率指标这个分工方式在课程设计级别的项目里算是规整的。Preprocessing.py 负责把 GTSRB 那套按类别文件夹存放的图片转成 CSV 索引TSRInput.py 再根据 CSV 去读图、做归一化和增广TSRCnn.py 只关心网络长什么样TSRTrain.py 把前面几个串起来跑训练TSREval.py 单独做验证。你如果要换数据集或者改网络结构改动的范围是可预期的不会牵一发动全身。2.2 GTSRB 数据集的目录结构与 CSV 索引的生成逻辑GTSRB 原始数据的组织方式是每个类别一个文件夹文件夹名就是类别 ID里面是一堆 .ppm 图片。这种结构人看着舒服但 PyTorch 的 DataLoader 不认所以 Preprocessing.py 的核心任务就是遍历目录、生成一份「图片路径 → 类别标签」的映射表。常见做法是用 os.walk 或者 glob 遍历把每张图的相对路径和它所属的文件夹名也就是标签写进 CSV。这里有个细节值得注意GTSRB 官方训练集有 43 个类别但每个类别的图片数量差异很大少的不到 200 张多的超过 2000 张。如果直接按原始分布训练模型会偏向样本多的类别。Preprocessing.py 里如果做了类别均衡或者至少统计了分布那说明作者是考虑过这个问题的。import os import csv import pandas as pd # GTSRB 原始数据根目录下面按类别 ID 分文件夹 data_root ./GTSRB/Training records [] for class_id in sorted(os.listdir(data_root)): class_dir os.path.join(data_root, class_id) if not os.path.isdir(class_dir): continue for img_name in os.listdir(class_dir): if img_name.endswith(.ppm): img_path os.path.join(class_dir, img_name) records.append({path: img_path, label: int(class_id)}) # 打乱后按 8:2 划分训练集和测试集 df pd.DataFrame(records).sample(frac1, random_state42).reset_index(dropTrue) split int(len(df) * 0.8) df[:split].to_csv(train_data.csv, indexFalse) df[split:].to_csv(test_data.csv, indexFalse) print(f训练集 {split} 张测试集 {len(df)-split} 张)这段代码的逻辑很直白先遍历所有类别文件夹收集路径和标签然后用 pandas 打乱并切分。random_state42是为了保证每次运行划分结果一致方便复现。frac1表示全量打乱。实际项目中如果要做交叉验证可以把这里的切分逻辑改成 KFold但课程设计级别用固定切分就够了。注意GTSRB 的测试集官方是单独提供的带标注的测试集有 12630 张。如果你用的是官方测试集就不需要从训练集里切分直接读官方测试标注文件即可。上面代码演示的是从训练集切分的做法适合官方测试标注不好找的情况。2.3 数据读取层的归一化与增广参数TSRInput.py 这个文件承担的是「把 CSV 里的路径变成模型能吃的张量」这个任务。GTSRB 图片尺寸不统一从 15×15 到 250×250 都有所以第一步肯定是 resize 到固定尺寸。这个项目里 CNN 的输入大概率是 32×32因为 GTSRB 的很多基准实验都用这个尺寸。from torch.utils.data import Dataset from PIL import Image import torchvision.transforms as T class GTSRBDataset(Dataset): def __init__(self, csv_file, img_size32, augmentFalse): self.df pd.read_csv(csv_file) self.img_size img_size # 基础变换resize 转张量 归一化 base [T.Resize((img_size, img_size)), T.ToTensor()] if augment: # 训练时加随机翻转和轻微颜色抖动 base [ T.Resize((img_size, img_size)), T.RandomHorizontalFlip(p0.3), T.ColorJitter(brightness0.2, contrast0.2), T.ToTensor() ] self.transform T.Compose(base) def __len__(self): return len(self.df) def __getitem__(self, idx): row self.df.iloc[idx] img Image.open(row[path]).convert(RGB) img self.transform(img) return img, row[label]这里有几个参数值得说清楚。img_size32是 GTSRB 上的经典选择再大当然可以但计算量会上去而且小图放大后信息增益有限。RandomHorizontalFlip(p0.3)的概率没有设成 0.5是因为交通标志里有左右不对称的类别比如某些指示箭头翻转太狠反而引入噪声。ColorJitter的 brightness 和 contrast 都控制在 0.2目的是模拟不同光照条件但幅度不能大否则标志的颜色特征会被破坏。提示归一化那一步我故意没写Normalize因为 GTSRB 的像素均值方差和 ImageNet 不一样直接用 ImageNet 的统计量会引入偏差。更稳妥的做法是先算一遍训练集的均值和方差再填进去。如果嫌麻烦至少把像素缩到 [0,1] 区间这对收敛速度的影响比想象中大。3. CNN 模型定义与训练循环从 TSRCnn.py 到 TSRTrain.py3.1 卷积层堆叠方式与感受野的匹配TSRCnn.py 定义的是整个项目的核心——网络结构。交通标志识别的特点是目标在图中占比大、背景相对干净、类别间差异集中在形状和颜色上。这意味着网络不需要太深的层次但卷积核的感受野要能覆盖整个标志区域。一个典型的做法是三层卷积加两层全连接。第一层用 5×5 或 3×3 的卷积核通道数从 3 升到 32第二层继续升到 64第三层升到 128。每层后面接 ReLU 和最大池化。这样经过三次池化后32×32 的输入变成 4×4感受野已经足够覆盖整张图。import torch.nn as nn class TSRCnn(nn.Module): def __init__(self, num_classes43): super().__init__() self.features nn.Sequential( # 输入 3×32×32 nn.Conv2d(3, 32, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), # 32×32 - 16×16 nn.Conv2d(32, 64, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), # 16×16 - 8×8 nn.Conv2d(64, 128, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), # 8×8 - 4×4 ) self.classifier nn.Sequential( nn.Flatten(), nn.Linear(128 * 4 * 4, 256), nn.ReLU(inplaceTrue), nn.Dropout(0.5), nn.Linear(256, num_classes) ) def forward(self, x): x self.features(x) return self.classifier(x)padding1配合kernel_size3保证卷积后空间尺寸不变这样池化层的输入输出关系就很清晰。Dropout(0.5)放在全连接层之间是为了防止过拟合GTSRB 训练集不到四万张不加 dropout 的话训练准确率很快冲到 99% 但验证集卡在 95% 左右上不去。num_classes43是 GTSRB 的标准类别数如果你只做其中几类改这个参数就行。3.2 训练循环里的学习率、批次大小与早停策略TSRTrain.py 是训练主循环。这部分代码的写法直接决定了你能不能跑出一个像样的结果。我见过太多课程设计项目在训练循环里翻车——要么学习率设太大导致 loss 震荡要么批次太小导致梯度噪声过高。import torch from torch.utils.data import DataLoader from TSRCnn import TSRCnn from TSRInput import GTSRBDataset device torch.device(cuda if torch.cuda.is_available() else cpu) model TSRCnn(num_classes43).to(device) train_set GTSRBDataset(train_data.csv, augmentTrue) val_set GTSRBDataset(test_data.csv, augmentFalse) train_loader DataLoader(train_set, batch_size64, shuffleTrue, num_workers2) val_loader DataLoader(val_set, batch_size64, shuffleFalse, num_workers2) criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.parameters(), lr1e-3) scheduler torch.optim.lr_scheduler.StepLR(optimizer, step_size10, gamma0.5) best_acc 0.0 for epoch in range(30): model.train() for imgs, labels in train_loader: imgs, labels imgs.to(device), labels.to(device) optimizer.zero_grad() loss criterion(model(imgs), labels) loss.backward() optimizer.step() # 每个 epoch 结束后在验证集上评估 model.eval() correct total 0 with torch.no_grad(): for imgs, labels in val_loader: imgs, labels imgs.to(device), labels.to(device) preds model(imgs).argmax(dim1) correct (preds labels).sum().item() total labels.size(0) acc correct / total scheduler.step() if acc best_acc: best_acc acc torch.save(model.state_dict(), best_model.pth) print(fEpoch {epoch1}, Val Acc: {acc:.4f}, Best: {best_acc:.4f})batch_size64是 32×32 小图上的常用值显存占用不大梯度也够稳。lr1e-3配合 Adam 是安全起点如果 loss 下降太慢可以调到 3e-3但再大就容易震荡。StepLR每 10 个 epoch 把学习率砍半目的是后期精细调整。早停策略这里用的是「保存最佳验证准确率对应的权重」没有做严格的 patience 早停但 30 个 epoch 对这个小网络来说已经够收敛了。注意num_workers2在 Windows 上如果报错改成 0 就行。这是 PyTorch 在 Windows 上的老问题跟代码逻辑无关。另外如果你没有 GPU把 device 改成 cpu 也能跑只是 30 个 epoch 可能要等一两个小时。3.3 评估脚本 TSREval.py 的指标输出与混淆矩阵训练完之后TSREval.py 负责给你一份能写进报告的结果。除了整体准确率交通标志识别这个任务特别需要看混淆矩阵因为有些类别长得太像了——比如限速 50 和限速 60或者各种禁止标志。from sklearn.metrics import classification_report, confusion_matrix import seaborn as sns import matplotlib.pyplot as plt model.load_state_dict(torch.load(best_model.pth)) model.eval() all_preds, all_labels [], [] with torch.no_grad(): for imgs, labels in val_loader: imgs imgs.to(device) preds model(imgs).argmax(dim1).cpu().numpy() all_preds.extend(preds) all_labels.extend(labels.numpy()) print(classification_report(all_labels, all_preds, digits4)) cm confusion_matrix(all_labels, all_preds) plt.figure(figsize(14, 12)) sns.heatmap(cm, annotFalse, cmapBlues) plt.xlabel(Predicted) plt.ylabel(True) plt.savefig(confusion_matrix.png, dpi150)classification_report会输出每个类别的 precision、recall 和 F1这比一个笼统的准确率有说服力得多。混淆矩阵热力图能让你一眼看出哪些类别容易被搞混。如果发现某两个类别的混淆特别严重可以考虑针对性补充数据或者调整网络最后一层的特征维度。4. 环境配置与运行链路从零把项目跑起来的完整步骤4.1 Python 环境与依赖库的版本选择这个项目依赖的东西不多核心就是 PyTorch、torchvision、pandas、Pillow、scikit-learn、matplotlib 和 seaborn。版本方面PyTorch 1.10 以上都能跑Python 3.8 到 3.10 都行。我一般建议用 conda 建一个独立环境避免和系统里的其他包打架。conda create -n tsr python3.9 -y conda activate tsr pip install torch torchvision pandas pillow scikit-learn matplotlib seaborn如果你用的是 pip 而不是 conda把前两行换成python -m venv tsr_env然后激活也行。PyTorch 的安装命令最好去官网查一下对应你 CUDA 版本的直接pip install torch默认装的是 CPU 版训练会慢很多。提示国内下载 PyTorch 有时候会很慢可以加-i https://pypi.tuna.tsinghua.edu.cn/simple用清华源。但注意有些镜像的 torch 版本更新不及时如果装不上就换回默认源。4.2 数据准备与目录结构的对齐GTSRB 数据集需要自己去官网下载下载下来是一个压缩包解压后得到 Training 和 Test 两个文件夹。Training 下面是 43 个类别文件夹Test 里面是一堆图片加一个 CSV 标注文件。你需要把 Training 的路径填到 Preprocessing.py 的data_root变量里然后运行它生成 train_data.csv 和 test_data.csv。# 假设目录结构如下 # TSR-master/ # Preprocessing.py # TSRTrain.py # ... # GTSRB/ # Training/ # 00000/ # 00001/ # ... # Test/ # *.ppm # GT-final_test.csv python Preprocessing.py python TSRTrain.py python TSREval.py运行顺序不能乱先预处理生成 CSV再训练最后评估。如果你改了网络结构重新跑 TSRTrain.py 就行不用重新生成 CSV。如果你换了数据集那 Preprocessing.py 要重跑。4.3 训练过程中的显存与日志观察训练跑起来之后终端会每个 epoch 打印一次验证准确率。你需要关注的是loss 有没有持续下降、验证准确率有没有在某个 epoch 之后停滞、显存占用有没有爆。如果显存不够把 batch_size 从 64 降到 32 或 16。如果验证准确率卡在某个值上不去先检查学习率是不是太大了再检查数据增广是不是太激进。# 在训练循环里加一行显存监控 if torch.cuda.is_available(): print(fGPU mem: {torch.cuda.memory_allocated()/1024**2:.1f} MB)这行代码放在每个 epoch 结束的地方能让你直观看到显存变化。如果显存一直涨不降可能是某个地方没有 detach 或者没有用torch.no_grad()检查评估部分有没有包在with torch.no_grad():里。5. 避坑与排查跑这个项目时最容易翻车的五个地方5.1 图片格式读取报错现象运行 TSRInput.py 时报UnidentifiedImageError或者cannot identify image file。原因GTSRB 原始图片是 .ppm 格式Pillow 对某些 .ppm 变体的支持不完整尤其是 P6 格式的二进制 ppm。解决在Image.open之前加一个格式转换或者用imageio库读图再转成 PIL Image。更省事的办法是在 Preprocessing.py 里就把 .ppm 转成 .png 或 .jpg后面读取就不会有问题。5.2 类别标签不连续导致 CrossEntropyLoss 报错现象训练时报IndexError: Target X is out of bounds。原因GTSRB 的类别文件夹名是 00000 到 00042但如果你手动删过某些类别或者用了自己整理的数据集标签可能不连续。CrossEntropyLoss 要求标签在 [0, num_classes-1] 范围内。解决在生成 CSV 的时候加一步标签重映射把实际出现的类别 ID 映射到连续的 0 到 N-1。或者直接用LabelEncoder处理。5.3 验证集准确率远低于训练集现象训练集准确率 99%验证集只有 85% 左右。原因过拟合。GTSRB 训练集不到四万张如果网络参数太多或者增广不够很容易过拟合。解决先加 Dropout再检查数据增广有没有开。如果还不够把全连接层的 256 降到 128或者加 L2 正则化。另一个容易被忽略的点是验证集的预处理必须和训练集一致除了增广部分。5.4 Windows 下 num_workers 报错现象RuntimeError: DataLoader worker (pid xxx) is killed by signal。原因Windows 上 PyTorch 的多进程 DataLoader 需要if __name__ __main__:保护否则会递归创建子进程。解决把训练代码包在if __name__ __main__:里面或者直接把num_workers设成 0。设成 0 的代价是数据加载变慢但对小数据集来说影响不大。5.5 模型保存后加载时报 key 不匹配现象RuntimeError: Error(s) in loading state_dict for TSRCnn: Missing key(s)。原因保存的时候用了torch.save(model, path)保存整个模型加载的时候用了model.load_state_dict()或者反过来。两种保存方式的加载方法不一样。解决统一用torch.save(model.state_dict(), path)保存加载时先实例化模型再load_state_dict。这样即使代码改了只要网络结构没变权重就能加载。6. 在 TSR-master 基础上做扩展换网络、换数据、提精度的实操思路把基础版本跑通之后如果你想让这个项目在答辩时更有亮点或者单纯想练手有几个方向可以试。第一个方向是换 backbone把 TSRCnn.py 里的三层卷积换成 ResNet-18 或 MobileNetV2 的预训练版本。做法很简单用 torchvision.models 加载预训练权重把最后的全连接层改成 43 类输出然后微调。但要注意 GTSRB 图片是 32×32而 ResNet 默认输入是 224×224你需要把图片 resize 到 224 再喂进去或者改 ResNet 第一层的 stride 和 kernel size 来适应小图。import torchvision.models as models def build_resnet18(num_classes43): model models.resnet18(pretrainedTrue) # 修改第一层卷积适应 32×32 输入 model.conv1 nn.Conv2d(3, 64, kernel_size3, stride1, padding1, biasFalse) model.maxpool nn.Identity() # 去掉 maxpool保留更多空间信息 model.fc nn.Linear(512, num_classes) return modelpretrainedTrue会下载 ImageNet 预训练权重第一次运行需要联网。model.maxpool nn.Identity()是为了不让 32×32 的图在早期就被池化得太小。model.fc替换成 43 类输出。这个改动之后验证准确率通常能比原始 CNN 高 1 到 2 个百分点但训练时间会变长。第二个方向是处理类别不均衡。GTSRB 里有些类别样本很少你可以用 WeightedRandomSampler 给少样本类别更高的采样权重或者在 loss 里给不同类别加权重。我一般会先跑一遍看混淆矩阵如果发现某几个类别的 recall 明显偏低再针对性处理。第三个方向是模型集成。把原始 CNN 和 ResNet 的预测结果做平均或者用投票法。这个在课程设计里算是加分项但实现起来不复杂就是把两个模型的 softmax 输出加起来再 argmax。从那以后我每次拿到一个新的图像分类项目都会先把数据集的类别分布画出来再决定要不要做均衡处理。这个习惯帮我省了很多事后调参的时间。希望帮到你。本文还有配套的精品资源点击获取