基于YOLOv5s改进的铁路信号灯小目标检测与部署实践

简介：面向铁路安全运输场景，这套深度学习实践资料围绕卷积神经网络（CNN）的铁路信号灯识别方法展开，适合图像识别入门者、计算机视觉方向学生及铁路智能监测相关研究人员。资源以普通铁路信号灯为研究对象，从CNN基础模型出发，逐步改进结构并优化学习率，解决信号灯位于地面与高处、多灯光混杂等条件下的颜色识别难题，为列车司机预判提供辅助思路。包内共141个文件，其中127张jpg图片构成红、黄灯等典型样本集，7个py脚本覆盖手写数字识别、颜色识别与信号灯实验代码，还附有tfrecords数据文件、png示意图、xls结果表及说明文档，整体压缩包约21.55MB，结构便于对照学习。目前已有289人浏览学习，适合需要从零搭建CNN颜色识别模型、开展实验对比与结果分析的读者，可据此理解图像特征提取与分类辨识的完整流程。

1. 项目概述与需求拆解

1.1 为什么铁路信号灯识别值得单独做一个项目

铁路信号灯这个东西，乍一看不就是红绿黄三个颜色嘛，能有多难？真正上手做过才知道，它跟马路上那种交通信号灯压根不是一个难度量级。列车运行速度高、制动距离长，信号灯没看清或者看晚了，后果非常严重。人眼在极端天气、强逆光、长时间驾驶疲劳的情况下，误判漏判的概率并不低，所以用视觉算法做辅助识别，在这个场景里有实实在在的需求。

更麻烦的是，铁路信号灯在画面里的占比通常非常小。举个例子，我用1080p的相机去拍几百米外的信号灯，灯体在图像里可能只有十几个像素，加上镜头畸变、振动模糊、晨雾夕阳这些干扰，传统图像处理那一套——颜色阈值分割、形状匹配——很容易就崩了。这正是我决定用深度学习来做识别而不是继续抠传统算法的主要原因。

除了技术层面，这个项目的工程价值也很明显。铁路巡检、车载辅助驾驶、道口监控，这些都是信号灯识别的落地场景。它能做到实时检测、持续运行，不依赖人工盯屏，一旦发现异常可以联动告警。对于做工业视觉、轨道交通相关的从业者来说，这类项目既是CV入门的一套经典练手内容，也是一条能写进简历的完整工程链路。

1.2 项目技术路线与选型逻辑

这个项目涉及的核心技术点集中在图像分类和目标检测两个方向。早先有一些方案把信号灯识别做成纯分类任务——先通过传统方法把灯的位置裁剪出来，再用CNN判断颜色和状态。这种方式的问题在于定位一旦不准确，分类再准也没用。所以我这版方案采用的是端到端的检测思路，把“灯在哪里”和“灯是什么颜色”这两个问题一起解决掉。

模型选型上，我对比过两个方向。两阶段检测器（比如Faster R-CNN）精度确实高，但推理速度在嵌入式设备上吃紧。单阶段检测器（比如YOLO系列）速度快，小目标召回率也能通过技巧补上来。综合权衡检测实时性和部署成本之后，我选了YOLOv5s作为基线模型，然后再针对信号灯这个小目标场景做专门的优化。具体到网络层面，我会在后面的章节把改动细节拆开讲。

数据层面，公共数据集里几乎没有专门针对铁路信号灯的高质量标注集，所以我一开始就把自建数据集当成核心工作来对待。这其实也是实际工程和学校作业最大的区别——模型结构可以复现，但数据才是决定项目上限的地方。

2. 数据集构建与预处理

2.1 数据采集的坑与经验

铁路信号灯的公开数据少，这是做这个项目首先要面对的现实。能拿到的原始数据来源大概有三类：一是铁路沿线的固定监控摄像头，二是车载前视相机录制，三是从一些公开的视频片段里抽帧。我自己主要用的是前两种，加起来采集了大约12000张图像。

采集过程有几个特别容易踩的坑。第一个坑是画面比例问题——监控摄像头往往是俯视角度，车载相机是平视角度，二者信号灯的形状特征差异很大。如果混着训练不处理，模型容易学到“角度特征”而不是“灯的特征”。我的做法是尽量让数据集里某一角度的图片占主导，然后用少量另一角度的图片做微调验证。

第二个坑是样本不平衡。绿灯和黄灯的样本数量远少于红灯——毕竟铁路信号里绿灯常见、黄灯相对少。如果不做处理，模型对少样本类别的召回率会非常难看。我的解决思路是尽量特意去多录一些黄灯出现的时间段，同时在后续做数据增强时对黄灯样本采用更高的增强倍率。

第三个坑涉及到标注一致性。信号灯目标小，标注框稍微画大一点画小一点，对训练结果影响都不小。我检查标完的框时就发现不少人习惯把灯外面的黑色背板也框进去，导致模型学了一堆背景信息。后来我重新统一了标注规范：框必须紧贴灯体发光区域，不包含背板和灯柱。

2.2 数据标注流程与格式转换

标注工具我用的是LabelImg，虽然老一些，但胜在稳定，支持Pascal VOC格式的XML输出。我的流程是先把所有抽帧图像按7:2:1划分成训练集、验证集、测试集，划分完之后再做标注——先划分再标注的好处是能避免同一条视频里前后帧的相似图像同时出现在训练集和测试集里，否则验证结果会虚高。

标注完成之后，要把VOC格式转换成YOLO训练需要的txt格式。每一行对应一个目标，格式是“类别id 中心点x 中心点y 框宽 框高”，所有坐标都归一化到0到1之间。这个转换本身不复杂，但写脚本时要注意坐标换算别出bug，尤其是从XML的左上右下坐标转成YOLO的中心宽高格式，容易搞混。

import os import xml.etree.ElementTree as ET def voc_to_yolo(xml_file, target_dir, class_names): tree = ET.parse(xml_file) root = tree.getroot() img_w = int(root.find('size/width').text) img_h = int(root.find('size/height').text) lines = [] for obj in root.iter('object'): cls = obj.find('name').text if cls not in class_names: continue cls_id = class_names.index(cls) box = obj.find('bndbox') xmin = float(box.find('xmin').text) ymin = float(box.find('ymin').text) xmax = float(box.find('xmax').text) ymax = float(box.find('ymax').text) x_center = (xmin + xmax) / 2.0 / img_w y_center = (ymin + ymax) / 2.0 / img_h w = (xmax - xmin) / img_w h = (ymax - ymin) / img_h lines.append(f"{cls_id} {x_center:.6f} {y_center:.6f} {w:.6f} {h:.6f}") txt_name = os.path.splitext(os.path.basename(xml_file))[0] + '.txt' with open(os.path.join(target_dir, txt_name), 'w') as f: f.write('\n'.join(lines))

类别这块我用的是red、green、yellow三个类别，先不做红色双灯、绿色双灯这种细分。灯芯数量、排列方式的分辨属于更细粒度的任务，放到后续迭代再处理。首批先保证三个基础颜色能稳下来，项目跑通了再叠加细节分类。

2.3 数据增强策略的针对性设计

信号灯目标小的特点决定了通用数据增强策略不能直接套用，得做针对性的调整。Mosaic增强在YOLOv5里是默认开的，它把四张图拼在一起训练，对小目标检测特别有效，因为拼图相当于把整体分辨率摊薄了，目标在画面中的相对尺寸变得更小，模型被迫去适应更严苛的检测条件。这个增强我保留了。

颜色抖动这块我做了专门强化。铁路场景的光照变化非常剧烈，大雾、雨天、逆光、夜晚都会影响灯体颜色表现。我在HSV空间做了较大范围的随机扰动，H通道扰动幅度控制在±0.02以内——因为色相一旦扰动过头，红灯的色相漂到橙色甚至黄色，就会给模型注入错误标注。S和V通道可以放宽到±0.3。

有一个增强操作我是刻意不用或者少用的：水平翻转。这个要结合信号灯的物理位置逻辑来考虑——铁路信号机是有固定排列规则的，比如有些地方红灯固定在左侧或上方。如果做水平翻转，等于强行制造了一批“红灯在右侧”的样本，跟真实物理场景分布不一致，反而会干扰模型学习位置先验。这种思路上我当时也查了不少资料，整体上倾向于保持原始空间关系。

3. CNN模型结构设计与改进

3.1 基线模型YOLOv5s的结构拆解

YOLOv5s的结构分三块：Backbone（骨干网络）、Neck（特征融合层）、Head（检测头）。Backbone用的是CSPDarknet，核心模块是C3结构——它把输入特征分成两路，一路经过若干Bottleneck模块提取深层特征，另一路直接跳过，最后在通道维度上拼接，这样在控制计算量的同时增强了梯度流动。Neck部分采用FPN+PAN结构，FPN自顶向下传语义信息，PAN自底向上传空间信息，两者结合确保不同尺度的特征图都能拿到足够信息。Head负责在三个不同尺寸的特征图上输出预测框和类别概率。

这个结构用在通用目标检测上很成熟，但直接拿来检测铁路信号灯有两个问题。第一，特征图下采样倍数太大。YOLOv5s的Backbone对输入图像做了32倍下采样，最后一层特征图上一个小目标的特征几乎被压没了。信号灯目标小，这一个问题直接决定了它容易漏检。第二，三个尺寸的检测头中，最大的特征图是80x80（在640输入下），本来应该负责检测小目标，但它的感受野和语义信息都是最浅层的，对小目标的表征能力有限——这是需要改进的地方。

3.2 针对信号灯小目标的三处修改

第一处修改是增加一个更高分辨率的检测头。原版YOLOv5s有三个检测头，对应8倍、16倍、32倍下采样的特征图。我额外加了一个4倍下采样的检测头，也就是输入640时输出160x160的特征图。这层特征图保留的空间信息最多，适合捕捉小尺寸信号灯。对应的，Neck部分也做了调整，让这条新检测头所在的支路能同时获得深层的语义信息，避免只靠浅层特征导致误检率升高。

第二处修改是在Backbone里引入注意力模块。我在C3模块后面接了一个轻量级的SE注意力模块。SE模块的思路是：先对特征图做全局平均池化得到每个通道的全局统计信息，再通过两个全连接层（先压缩再扩展）学习每个通道的重要性权重，最后把权重乘回原特征图。相当于告诉模型“什么特征重要就强调什么”。对信号灯这种目标小、但颜色和形状特征明显的场景，通道注意力能有效增强语义响应。

import torch.nn as nn class SEBlock(nn.Module): def __init__(self, channels, reduction=16): super().__init__() self.squeeze = nn.AdaptiveAvgPool2d(1) self.excitation = nn.Sequential( nn.Linear(channels, channels // reduction), nn.ReLU(inplace=True), nn.Linear(channels // reduction, channels), nn.Sigmoid() ) def forward(self, x): b, c, _, _ = x.size() y = self.squeeze(x).view(b, c) y = self.excitation(y).view(b, c, 1, 1) return x * y.expand_as(x)

第三处修改是替换激活函数。原版YOLOv5s在C3模块里用的是SiLU激活函数，我在深层网络部分尝试替换成了ReLU，并结合权重初始化调整做了对比实验。这里说下结论：ReLU和SiLU在信号灯检测任务上精度差异不大（约0.3%mAP），但ReLU在推理时对硬件更友好，某些嵌入式设备的加速库对ReLU的支持更完善。如果只做实验不搞部署，保持原版SiLU完全没问题。

3.3 模型轻量化探索记录

我还尝试过用MobileNetV3替换Backbone做轻量化。MobileNetV3用深度可分离卷积替代普通卷积，参数量和计算量大幅下降，但实验效果显示，在信号灯数据集上，替换后的mAP下降了约4.5个点，主要是小目标召回率掉了。后续我有做过分析，小目标检测对浅层空间分辨率敏感，MobileNetV3为了轻量把浅层通道数压得比较低，不利于保留小目标的细节特征。

有一个方向值得继续试：FPN结构加双向融合（BiFPN），它在相邻层之间增加了跨尺度连接。原版YOLOv5s的PAN虽然也是双向，但BiFPN会给不同输入特征学习权重，理论上对小目标的特征融合更友好。不过BiFPN带来的额外计算量跟精度收益之间的平衡，在当前数据集上还需要更多实验验证。这个我在文末会再提一句。

4. 训练过程与关键参数

4.1 训练环境与超参数配置

训练环境这块，我用的是一块RTX 3090显卡，显存24GB。软件栈是Python 3.9、PyTorch 1.12、CUDA 11.6。这里有个经验可以分享：PyTorch版本和CUDA版本一定要匹配，否则编译某些自定义算子的时候会报莫名其妙的错，排查起来浪费时间。

核心超参配置如下：

这里特别注意SGD和Adam的选择。很多人习惯直接用Adam，但YOLOv5的默认配置SGD在检测任务上往往更稳，训练到后期不容易出现loss震荡的问题。Adam收敛快，但最终精度通常不如调好的SGD。这个差异在信号灯数据集上实测差了接近1个mAP点。

4.2 训练曲线怎么判断是否正常

训练要盯的不只是loss降没降，更关键的是验证集mAP有没有跟着涨。我之前有过一次训练，loss一路降到很低，但mAP始终上不去——后来发现是训练集和验证集之间数据分布出现了明显差异，训练集里大量是同源视频抽帧，特征太相似了。

正常训练的loss曲线特征是：前10轮下降非常快，从初始的0.1以上迅速降到0.05以下；之后下降变缓，在60轮左右进入平台期。如果loss降到平台期之后还在缓慢下降，说明模型还在学有效特征，不必着急提前停止。mAP曲线和loss曲线大致呈镜像关系，但会有小幅滞后，这是正常的。

学习率的调整策略我直接用了YOLOv5自带的余弦退火，它在每个epoch都会调整学习率，让学习率从初始值逐渐降到接近0。相比固定学习率，这种方式在训练后期能有效避免参数在最优解附近震荡，稳定性和最终精度都有提升。

4.3 训练后模型表现数据

120轮训练完成后，我在测试集上的模型表现如下：

可以看到yellow的召回率显著低于red和green，这跟样本量少直接相关。虽然做了类别加权和增强策略，但数据量本身不够这个先天因素不是完全能靠技术弥补的。后续要提高黄灯召回率，最有效的还是继续采集黄灯样本。

5. 部署实践与实时推理优化

5.1 从PyTorch模型到TensorRT部署

训练好的模型不能直接拿到现场跑，PyTorch模型在GPU上推理一张640x640图像需要几十毫秒，听起来还行，但如果跑在车载嵌入式设备上，这张模型动辄几百毫秒，而且对显存占用太高。所以部署这一步必须做模型转换和推理加速。

我的部署流程是：PyTorch模型转ONNX，再做TensorRT下的FP16量化。ONNX转换这一步要特别注意，YOLOv5的模型里有不少操作是PyTorch自定义的，转ONNX时可能会产生多余的reshape节点，影响后面TensorRT的执行效率。如果遇到这个问题，可以尝试把模型输入输出的attribute固定住，减少动态维度。

转换成TensorRT之后，我在同一张3090上做了推理速度对比：PyTorch原生推理约22ms/帧，ONNX Runtime约15ms/帧，TensorRT FP16量化后约6ms/帧。提升幅度相当可观。FP16量化带来的精度损失很小，实测mAP只掉了约0.3个百分点，对信号灯识别这个任务来说完全可接受。

5.2 若干实用推理优化技巧

推理阶段还有一个容易被忽视的点：图像预处理的方式。检测网络通常用letterbox方式将输入图像等比缩放，不足部分补灰边，避免目标变形。但letterbox会改变目标在图像中的实际位置坐标，所以推理时要把预测框坐标反向映射回原图坐标系，这个映射公式一定要写对，否则部署出去框位置会偏移。

另外，滑动窗口做高分辨率图像推理也是一种方案。铁路监控摄像头输出的图像往往不止1080p，直接resize到640会丢失大量小目标信息。这种情况下可以把原图切块，每块独立推理后再合并NMS。代价是推理时间成倍增加，只能用在准实时场景。

电力或硬件资源受限的场景下，还可以考虑用OpenVINO做CPU推理。Intel的核显CPU跑YOLOv5s 640输入可以做到30ms左右一帧，虽然比TensorRT慢，但胜在不需要独立显卡，适合低成本改造既有监控系统。这个方向我在项目里做了一部分验证，后面有时间可以单独展开写一篇。

6. 常见问题与排查记录

6.1 漏检红灯问题

训练完成之后测试，发现一个非常要命的场景问题：在远距离、大逆光的条件下，红灯经常漏检。排查思路分两步走。

第一步，我先看是不是训练数据里缺少这种“逆光远距离红灯”的样本。统计结果证实了猜测：这类样本占比不到2%。解决办法是专门针对这种场景做数据扩充——手动收集逆光时段的数据并补充到训练集里。

第二步，在模型层面调低类别置信度阈值。YOLO系列默认的conf_thres是0.25，对于信号灯检测这种宁可误检也不能漏检的场景，我把阈值降到了0.1。代价是误检数量确实变多了，但配合后处理的“同一位置连续多帧确认逻辑”，可以把误检率压到可接受水平。

6.2 夜间灯光泛光导致误检

夜间的信号灯存在一个特殊现象——发光灯体会出现泛光，光线在镜头里晕开一大片，跟灯体实际物理尺寸严重不符。模型一开始会把整个泛光区域都识别成灯，预测框比真实灯体大好几倍。

解决这个问题的最直接方案是在数据增强阶段增加模拟泛光的数据变换。我在HSV扰动之外额外加了一个操作：对部分训练样本做高斯模糊叠加，同时随机放大亮度通道的值，模拟夜间灯光扩散效果。经过这样处理后，模型对泛光区域的预测框收敛显著，明显更贴合真实灯体位置。

另外一个排查技巧是在推理阶段对预测框做宽高比过滤。铁路信号灯基本上是圆形或接近圆形，宽高比接近1。如果预测框宽高比超过2.0以上，基本可以断定是误检，直接丢掉即可。

6.3 雨雾天气识别率下降的实测体验

雨雾天气是信号灯识别最难对付的场景。雾气导致灯体颜色衰减严重，红灯看起来偏灰粉，黄灯看起来偏白。我在雨雾样本上的实测mAP掉到了70%左右，跟晴天场景差了20多个点。

这个问题的根治方案是采集大量雨雾天气的样本加入训练，但实际条件不可能等下雨再去采集。退而求其次的替代方案是做“雾化增强”合成数据——用暗通道先验原理对晴天样本做合成雾化模拟，按不同浓度增强后送入训练。我的实验数据显示，加入雾化增强后，雨雾测试集mAP从70%提升到了76%，有效但有限。真正想做得更好，还是需要真实的恶劣天气数据。

7. 常见问题速查表

8. 写在最后的实操建议

这个项目从数据采集到模型部署，完整走下来花了三个多月。如果让我做一次复盘，有一个建议最值得分享给想复现这个项目的朋友：第一，把时间花在数据上永远比花在调参上划算，先认真分析和补齐数据分布，再考虑改网络结构；第二，小目标检测的关键突破口往往在网络输入分辨率和检测头的设计上，与其在Attention上花大量功夫，不如先把这两个基础维度的效果做到位。

最后再补充一个容易被人忽视的细节。训练和推理时对图像做的预处理必须保持一致。比如训练时letterbox填充用的是灰色像素值114，推理时如果用了0或者255去填充，模型效果会显著变差。这种细节问题很难通过看日志发现，但排查起来又特别简单——对比训练和推理代码中的预处理函数，最好抽出来封装成同一个模块复用。这个习惯帮我省了很多排查问题的时间，也一并分享给各位。

本文还有配套的精品资源，点击获取