简介目标检测是计算机视觉的核心任务之一旨在从图像或视频中定位并识别特定对象。YOLOv5作为单阶段检测算法的代表通过回归方式直接输出目标位置与类别在保持高精度的同时实现了极快的推理速度尤其适用于需要实时响应的安全监管场景。PyTorch则凭借动态计算图和成熟的CUDA支持为模型训练与调试提供了极大便利。基于这两项技术构建的头盔检测系统可自动识别监控画面中人员是否佩戴安全帽一旦发现违规立即告警广泛应用于工地、园区、工厂等安全生产领域。本文从环境配置、数据标注、模型训练到推理部署完整呈现了构建此类系统的全过程并分享了解决漏检、误检及显存不足等实际问题的经验为开发者提供了一套可复现的工程实践方案。 去年接了一个工地的安全监管需求靠人盯监控看工人有没有戴安全帽眼睛根本看不过来。后来就把这个需求落地成了基于YOLOv5和PyTorch的深度学习头盔检测系统能做到对图片、视频、实时监控画面中的头盔佩戴情况进行自动识别检测到未佩戴立刻框出来并告警。整套流程从环境搭建、数据标注、模型训练到最终部署推理我完整跑了一遍这篇文章就把里面的关键细节和踩过的坑一次性讲清楚。适合想用深度学习做目标检测但还停留在看过教程、没跑过完整项目阶段的人也适合准备在工地、园区、工厂等场景落地安全帽检测的开发者参考。1. 项目整体设计与技术选型思路1.1 检测算法选型为什么锁定YOLOv5头盔检测本质上是一个目标检测任务核心要解决两个问题第一把画面里的人找出来第二判断这个人头部区域有没有佩戴头盔。市面上可选的目标检测算法不少两阶段的Faster R-CNN、单阶段的SSD和YOLO系列我最终选了YOLOv5原因很直接。YOLO系列发展到v5这代工程化程度已经非常高。相比Faster R-CNN那种先提候选区域、再逐个分类回归的两阶段方案YOLOv5把检测当成一个回归问题一次前向推理直接输出所有目标的位置和类别速度优势非常明显。对于头盔检测这种需要实时响应的场景推理速度是硬指标一秒钟处理几十帧画面才有实用价值。另外YOLOv5的代码结构在同期开源项目里属于最清晰的那一档。模型定义、训练流程、数据加载、推理脚本全部模块化改配置就能换模型尺寸和数据路径这对我们这种要把项目真正落地、后期还要根据不同摄像头场景反复调优的团队来说太重要了。你想换YOLOv5s还是YOLOv5m改一行参数就行不用动代码。还有一个关键因素预训练权重。YOLOv5官方提供了在COCO数据集上预训练好的权重虽然COCO里没有头盔这个类别但模型已经学到了通用的特征提取能力——边缘、纹理、形状这些底层特征我们只需要在它的基础上用自己的头盔数据集做微调训练时间和数据量需求都能大幅降低。这跟让一个见过几万种物体的人去学认一个新物种比让婴儿从零开始学要快得多是一个道理。1.2 深度学习框架选型PyTorch为什么顺手框架选了PyTorch基本没纠结。一个项目里框架的选择往往决定了开发效率的上限PyTorch最大的优势是动态计算图模型结构可以在运行时改变调试的时候你能随时打印每一层的输出形状和梯度信息遇到问题可以非常直观地定位。举个例子训练过程中如果某一层的输出维度对不上PyTorch会直接在报错信息里告诉你张量的形状和预期值配合print或者pdb断点几乎不需要猜。相比之下静态图框架是先构图、后执行调试时候的反馈链路会绕很多。对深度学习的初学者和中等水平开发者来说PyTorch的学习曲线也友好得多社区里能找到的海量教程和开源代码大部分都是PyTorch版本遇到问题一搜基本都有答案。还有一点是硬件生态。我们当时训练用的就是NVIDIA的GPUPyTorch对CUDA的支持非常成熟从显卡驱动到CUDA Toolkit再到PyTorch版本官方文档给出了完整的对应关系。训练时一个model.cuda()就能把网络搬到GPU上nvidia-smi可以实时查看显存占用非常顺手。为什么用PyTorch而不用TensorFlow这个问题我自己的答案很朴素YOLOv5官方源码本身就是基于PyTorch实现的选框架的第一个原则是跟随项目生态能直接跑通、能快速迭代比什么都重要。框架本身只是个工具好用、社区活跃、坑少就是最大的竞争力。1.3 系统整体架构与工作流程整个头盔检测系统可以拆成离线训练和在线推理两条链路。离线训练链路是数据采集和标注、数据集划分、模型训练、效果评估产出一个训练好的权重文件。在线推理链路是读取图片或视频流、送入模型做前向推理、后处理非极大值抑制、在画面上绘制检测框和置信度、根据检测结果触发告警逻辑。两条链路都围绕同一套模型文件运转这也是我比较推荐的落地方式。训练和推理解耦以后你可以在训练阶段随意折腾模型结构和超参数一旦权重固定下来推理端就可以做成一个轻量服务比如用Flask对模型做一层封装再对接监控系统的告警模块。我们当时的告警逻辑做了个简单的业务规则检测到人但没检测到头盔且人的检测框与头盔检测框没有重叠就判定为未佩戴触发记录和提醒。这里有个细节直接用人和头盔两个类别的检测结果做逻辑判断比单独训练一个佩戴/未佩戴分类模型要灵活得多因为画面里人很多逐个分类会增加漏检风险而目标检测可以同时输出所有人的位置和属性规则也好调整。2. 环境搭建从零配置PyTorch与YOLOv52.1 深度学习环境的三层结构环境搭建往往是最劝退新手的一步很多人在这一步就被各种版本冲突搞得心态崩溃。我把深度学习环境拆成三个层级来理解会清晰很多。最底层是驱动和CUDA。显卡驱动是操作系统和GPU硬件之间的桥梁CUDA Toolkit是NVIDIA提供的并行计算平台PyTorch在GPU上跑底层就是调用CUDA来操作GPU算力。往上走是Python环境我强烈建议用Anaconda来管理每个项目创建独立的虚拟环境互不干扰。最上层才是PyTorch框架本身。一个常见的误解是CUDA装得越新越好。其实不是关键是驱动版本要支持你选的CUDA版本而PyTorch版本又要对应匹配某个CUDA版本。这块我踩过一次坑有台机器装了很新的驱动我直接装了最新版CUDA结果那个版本的PyTorch还不支持又得重来一遍。后来我学乖了一律按PyTorch官方安装命令里指定的CUDA版本来驱动驱动的安装。用nvidia-smi查看显卡驱动版本右上角的CUDA Version表示该驱动最高支持的CUDA版本只要你的CUDA Toolkit版本不高于这个数字基本就没问题。然后去PyTorch官网选对应组合官网会给出安装命令直接复制执行即可。2.2 PyTorch GPU版本的安装与验证我们当时的安装路线是在Ubuntu 22.04上配置的命令大概是这样的conda create -n helmet python3.9 conda activate helmet pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118这里有两个点要特别说明。第一创建环境时Python版本选了3.9不是最新的3.11或者3.12原因是YOLOv5的依赖库尤其是numpy、opencv-python这些在3.9下兼容性最稳。深度学习项目的依赖往往很敏感Python小版本升级可能导致某个库没有对应的预编译包所以建议用版本适中、生态成熟的Python环境。第二PyTorch的安装命令里cu118表示CUDA 11.8版本。为什么选11.8而不是更高的12.x因为当时YOLOv5的官方测试环境就是11.8社区反馈这个组合下各种算子都稳定没必要追求最新。安装完成之后一定要验证GPU能不能正常调用我一般会跑这三行import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))如果输出True和显卡型号说明PyTorch已经能正常调用GPU。这一步验证不能省很多人装完就急着训练结果模型一直在CPU上跑速度慢得离谱还以为是代码问题。2.3 获取YOLOv5源码与安装依赖PyTorch环境就绪后接着获取YOLOv5源码。这里我建议直接用Git拉官方仓库后续要更新代码或者切换版本也方便git clone https://github.com/ultralytics/yolov5 cd yolov5 pip install -r requirements.txtrequirements.txt里包含了YOLOv5运行所需的全部依赖包括numpy、opencv-python、matplotlib、seaborn这些。有些教程会让你一条条手动安装没必要一条命令搞定。装依赖的时候有个小坑默认PyPI源在国内拉取很慢有些大包比如opencv-python动辄几十兆容易超时。我的做法是直接用国内镜像源把pip下载地址切到清华源或者阿里源速度能快一个数量级pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple环境装好后可以直接跑一下官方自带的推理脚本用预训练权重yolov5s.pt去检测自带的测试图片确认整个链路是通的再开始准备自己的数据。这一步相当于验收环境确保问题不堆积到后续环节。3. 头盔数据集准备与标注3.1 数据集来源与规模规划深度学习是数据喂出来的模型效果好不好的下限由数据质量决定。头盔检测的数据集有两个主要来源一是公开数据集二是自己采集标注。我们在项目初期先用公开数据集把整套训练流程跑通后期再针对实际场景补充自采数据。公开数据集里有几个可以重点关注SCUT-HEAD这是华南理工大学开源的头部检测数据集包含超过4000张标注好的头部图像还有其他研究者发布的建筑工地安全帽佩戴检测数据集标注了head和helmet两个类别。我当时的做法是把几个公开数据集合到一起做了一个大约8000张图的初始数据集类别就两类helmet和head。这里有个很关键的经验不要贪多求全。头盔检测的类别设计越简单越好我见过有人把类别分得很细什么红色安全帽白色安全帽蓝色安全帽结果每个类别的样本量都不够模型训练出来mAP惨不忍睹。头盔就是头盔检测目标就是有没有戴类别分多了只会增加模型的负担。数据集规模上头盔检测这类任务两张类别每类3000到5000个标注实例已经能训练出一个可用的模型。如果完全从零开始训练没有预训练权重可能需要更多数据但我们用了YOLOv5的COCO预训练权重做迁移学习数据需求量就小很多。3.2 数据标注labelImg标注头盔类别数据标注是项目里最枯燥但最重要的一环直接决定模型精度的上限。我们用的标注工具是labelImg一个基于Python的图形化标注工具体积小、操作简单、跨平台标注结果直接保存为YOLO格式的txt文件。安装labelImg很简单conda install pyqt5 pip install labelImg labelImg打开工具后用Open Dir打开图片文件夹然后通过Change Save Dir设置标注文件保存路径再在左侧选择类别最常用的快捷键W是创建标注框A和D切换上一张/下一张图CtrlS保存。标注的核心原则是框得准、不要漏。画框的时候框体要紧贴目标边缘稍微留一点点边距可以但不要框得太大把背景包进去也别框得太小截断了目标的一部分。一个容易忽略的细节是画面里非常小、非常模糊的目标也要尽量标注因为现实场景中摄像头距离远的时候目标就是很小的一块如果训练数据里全是清晰的大目标模型在真实场景下的召回率会很低。3.3 数据集目录组织与data.yaml配置标注完成后数据集要按YOLOv5要求的目录结构组织起来。YOLOv5的数据加载逻辑是images目录放图片labels目录放标注文件图片和标注文件通过文件名一一对应后缀不同.jpg和.txt名称相同。目录结构如下helmet_dataset/ ├── images/ │ ├── train/ # 训练集图片 │ └── val/ # 验证集图片 └── labels/ ├── train/ # train图片对应的标注txt └── val/ # val图片对应的标注txt训练集和验证集的划分我习惯按9:1或者8:2来但划分之前必须做一步去重。因为很多公开数据集的图片往往是从视频里逐帧截取的连续帧之间相似度极高如果不做去重模型会在验证集上表现很好到了真实场景就打回原形这就是过拟合的典型表现。我的做法是用脚本计算图片的感知哈希把相似度高的图片归到同一个集合里确保训练集和验证集的内容足够独立。标注文件的内容是YOLO格式的纯文本每一行代表一个目标对象格式为类别id x_center y_center width height四个坐标值都是归一化到0到1之间的比例值。比如0 0.5125 0.3452 0.2341 0.1876表示一个类别为0的目标中心点在图片的横向51.25%、纵向34.52%的位置宽高分别占整图的23.41%和18.76%。最后配置一个data.yaml文件内容很简单train: /path/to/helmet_dataset/images/train val: /path/to/helmet_dataset/images/val nc: 2 names: [helmet, head]nc是类别数names是类别名称列表顺序必须跟标注文件里的类别id一致。这个data.yaml在训练时用--data参数指定YOLOv5会自己读取。4. 模型训练参数理解与效果调优4.1 从卷积到检测YOLOv5核心结构速览训练之前哪怕不打算改模型结构也建议把YOLOv5的网络结构过一遍否则调参的时候会一头雾水。YOLOv5整体可以分成三块Backbone、Neck和Head。Backbone负责提取图像特征使用的是CSPDarknet结构。它的特点是引入了Cross Stage Partial结构把特征图分成两部分一部分走卷积层做特征提取另一部分直接连接到后面既减少了计算量又解决了深层网络梯度消失的问题。你可以把它理解为一条高速公路信息既能走铺满细节的小路也能走无障碍的直达通道。Neck部分用的是PANet结构它的作用是融合不同尺度的特征。头盔有大有小——近处的头盔占半个屏幕远处的头盔可能就十几个像素。PANet通过自顶向下和自底向上的双向路径把浅层的细节信息和深层的语义信息融合在一起让模型对小目标更敏感。Head部分输出三个不同尺寸的检测结果分别对应大、中、小目标。这也是YOLO系列的经典设计每个网格单元预测若干个锚框每个锚框输出位置偏移量、宽高、置信度和类别概率。至于深度学习的池化这个知识点在YOLOv5中也扮演了角色。SPPSpatial Pyramid Pooling模块在Backbone的最后一层使用不同池化核尺寸对特征图做池化再把结果拼接起来目的是增强模型对不同尺度目标的感受野。简单说池化就是压缩信息、扩大视野让网络从更大范围去理解一个目标是什么。4.2 训练超参数每一行配置背后的逻辑训练头盔检测模型最常用的是YOLOv5s这个尺寸的模型。yolov5s是官方提供的五种尺寸中最小的速度最快精度也不错。对头盔检测这种目标特征相对简单的任务用yolov5l甚至yolov5x属于杀鸡用牛刀训练时间长、推理速度慢精度提升却很有限。启动训练的命令cd yolov5 python train.py --data helmet.yaml --cfg yolov5s.yaml --weights yolov5s.pt --epochs 100 --batch-size 16 --img 640几个关键超参数我逐个说一下。--epochs是训练轮数。100轮是一个比较合理的起点YOLOv5训练过程中会自动保存效果最好的权重所以轮数多一点不用担心过拟合配合早停机制就行。我见过有人只训练30轮就停下来用效果好不到哪去因为模型还没充分收敛。--batch-size是批次大小。这个参数受限于显存显存不够就调小一点。但注意batch-size大小会直接影响训练效果太大的batch会让模型收敛不稳定太小的batch又会导致梯度估计噪声大。--img是输入图片的尺寸。默认640x640这个尺寸对大多数场景够用。头盔检测如果主要看监控画面目标通常比较小可以考虑用更大尺寸比如960或者1280来训练对小目标检测有提升但对应的显存消耗和推理耗时也会增加。还有一组藏在data/hyps/hyp.scratch-low.yaml里的超参数它们是深度学习训练中容易被忽略但很关键的细节。比如lr0是初始学习率momentum是动量weight_decay是权重衰减系数。这些参数YOLOv5官方已经调得比较合理新手不建议动等你跑完一轮训练对损失下降曲线有感觉了再来调学习率也不迟。4.3 启动训练与过程监控训练跑起来后不是干等着要会看训练日志。YOLOv5每个epoch结束会输出一组指标包括box_loss、obj_loss、cls_loss和mAP。box_loss代表预测框位置和真实框的误差这个值应该在训练中逐渐下降obj_loss代表目标置信度误差也需要下降cls_loss是分类误差对头盔检测来说前期会降得比较快因为两类区分度很高。mAP0.5是最直观的指标表示IoU阈值0.5下的平均精度均值头盔检测任务一般能训练到0.85以上就算非常可用的模型了。训练过程中有一个现象第一次遇到可能会慌损失值在初始几个epoch会跳来跳去甚至短暂上升这是正常的。因为迁移学习初期模型在快速调整底层特征损失波动不代表训练失败。我的经验是前20轮不用频繁关注指标到30轮以后再看趋势才有意义。显存不足的情况在训练中很常见报错通常是CUDA out of memory。最简单的处理是调小--batch-size从16降到8或者4如果还是不足就换小尺寸的模型。另外可以在训练命令中加上--cache参数把数据预加载到内存中能减少GPU和CPU之间的数据搬运等待但并不降低显存占用。4.4 训练结果评估mAP、PR曲线怎么读训练结束后YOLOv5会在runs/train/目录下生成实验记录里面有results.png、confusion_matrix.png、PR_curve.png等图表这些图对一个深度学习项目来说就是体检报告。PR_curve.png是精度-召回率曲线。精度Precision衡量的是模型预测为头盔的框里有多少真正的头盔召回率Recall衡量的是真实头盔中有多少被模型找出来了。这两个指标天然矛盾鱼和熊掌不可兼得。PR曲线与坐标轴围成的面积就是AP值面积越大表示模型效果越好。头盔检测这种安全相关的场景我个人更看重召回率——漏检一个没戴头盔的人比多误报几次更严重。所以后期调参时可以适当降低置信度阈值提高召回率。confusion_matrix.png是混淆矩阵能直观看出模型在哪两类之间混淆。对头盔检测来说最常见的混淆就是头盔和头部互相误判比如把戴浅色帽子的人的头误判为头盔。如果混淆矩阵里这两类的交叉区域比较大说明训练数据里头盔和头部的外观特征区分度不够需要补充更多难例数据。还有一个容易被忽视的评估指标是各类别的AP单独值。YOLOv5训练日志里会按类别分别输出AP如果helmet类AP很高但head类AP很低说明数据集中head类的样本不够或者标注质量有问题需要针对性补充。5. 实时检测从模型到应用5.1 使用训练好的权重进行图片和视频推理训练完成后最佳权重保存在runs/train/exp/weights/best.pt。使用这个权重做推理很简单python detect.py --weights runs/train/exp/weights/best.pt --source data/images --conf-thres 0.5--source既可以是图片文件夹也可以是一个视频文件路径甚至是摄像头设备的索引号比如--source 0表示读取本机第一个摄像头。推理时的一个关键调整参数是--conf-thres即置信度阈值。这个值的设置直接影响系统的误报和漏报阈值设高比如0.7模型只输出有把握的结果误报少但容易漏检阈值设低比如0.3漏检少但误报多。我建议头盔检测场景设0.4到0.5之间兼顾两种情况。detect.py跑完后结果图片和视频会保存在runs/detect/目录下检测框会标注类别名称和置信度。这里有个小技巧可以给不同类别设置不同颜色和线宽让未佩戴头盔的人脸框显示为红色醒目提示而已佩戴的框显示为绿色这对接实际业务体验很重要。5.2 实时视频流检测与部署detect.py默认方式是对视频逐帧检测然后保存结果但真正的实时检测系统需要把模型集成到自己的代码流程里。核心思路是先加载模型然后循环读取视频帧每帧做一次前向推理。一个简化的检测循环框架import torch import cv2 model torch.hub.load(ultralytics/yolov5, custom, pathruns/train/exp/weights/best.pt, force_reloadTrue) model.conf 0.5 model.iou 0.45 cap cv2.VideoCapture(0) # 读摄像头或视频流 while True: ret, frame cap.read() if not ret: break results model(frame) # results.pandas().xyxy[0] 包含所有检测框的坐标、类别和置信度 # 在这里编写告警逻辑有人但无头盔 - 发送告警 rendered results.render()[0] cv2.imshow(Helmet Detection, rendered) if cv2.waitKey(1) 0xFF ord(q): breaktorch.hub.load是PyTorch提供的一个非常方便的功能可以从本地加载自定义训练好的模型不需要手动构造网络结构。模型加载后通过设置model.conf和model.iou控制置信度阈值和NMS的IoU阈值这也是在线部署时最常用到的两个后处理参数。实时检测的实测性能需要心里有数。YOLOv5s在NVIDIA GTX 1660 Ti显卡上640x640输入推理速度大概在3到5毫秒一帧也就是每秒可以处理200到300帧实时性完全不是瓶颈。真正的瓶颈往往在视频解码和显示环节cv2.VideoCapture读RTSP流时延迟可能达到几百毫秒这是摄像头侧和网络侧的问题跟模型无关。5.3 模型导出与边缘设备部署模型训练好之后如果只在自己电脑上跑演示没问题但实际项目大多要部署到边缘设备或者服务器上。YOLOv5官方支持将模型导出为多种格式ONNX、TensorRT、CoreML、OpenVINO。导出ONNX格式的命令python export.py --weights runs/train/exp/weights/best.pt --include onnx --img 640导出ONNX的意义在于跨平台部署。ONNX是一种开放的模型表示格式相当于深度学习的通用语言很多推理引擎和边缘设备芯片比如RK3568、Jetson系列都支持加载ONNX模型。我们在项目中用了瑞芯微的RK3568平台做边缘部署流程是PyTorch训练 - 导出ONNX - 转换成RKNN格式 - 在板子上跑推理。导出时有个细节需要留意--img 640要和训练时的输入尺寸一致否则导出后模型可能因为输入尺寸不匹配而报错或者精度下降。另外如果想在TensorRT上部署建议先导出ONNX再转TensorRT引擎而不是直接导出因为兼容性更好。Jetson设备上部署PyTorch模型是另一个常见需求。注意Jetson平台使用的JetPack版本决定了支持的PyTorch版本比如JetPack 6.x一般对应PyTorch 2.x安装时不要用桌面版通用的pip命令而是参考NVIDIA官方针对Jetson提供的预编译wheel包这样能避免很多兼容性问题。6. 常见问题与排查技巧实录6.1 训练不收敛的排查思路训练过程中最让人崩溃的就是loss不降或者降着降着突然变成NaN非数值然后训练崩了。这个问题我遇到过不止一次原因是多方面的排查顺序建议按照数据 - 超参数 - 环境来。先说数据层面。检查标注文件里有没有异常的框比如坐标值超过1、宽高为0这种脏数据。YOLO格式要求归一化坐标在0到1之间如果标注工具或脚本出了问题某个坐标写成了5.7模型训练时梯度就会爆炸。我有一个脚本专门扫描数据集里的非法标注跑一遍就能发现并剔除问题样本。其次是学习率。如果loss前几轮就出现NaN大概率是学习率太大。YOLOv5的默认初始学习率是0.01如果你的数据集很小可以把hyp.scratch-low.yaml里的lr0调到0.005再试。学习率调小了一般能解决代价是收敛速度变慢多跑几轮就行。环境层面的问题比较隐蔽。有一种情况是GPU驱动和CUDA版本不匹配导致某些算子计算异常表面上看loss一切正常但训练到中途就开始跳NaN。检查方法是先用官方模型跑一个小的训练任务比如只训练5轮、只用100张图如果官方模型也崩那问题就不在你自己的数据和代码而是环境层面回退CUDA或者PyTorch版本重试。6.2 显存不足的解决路径CUDA out of memory是新手最常遇到的报错看到一堆红色报错信息容易慌其实解决路径很清晰。第一步把batch-size调小。训练命令中--batch-size 16改成8如果还爆就改成4。batch-size是显存占用的主要因素调小它效果立竿见影。第二步是换小模型YOLOv5s已经是小模型了如果用的yolov5m或者更大换回yolov5s。第三步是减小输入图片尺寸--img 640改成--img 416显存占用会下降不少但检测精度也会相应降低。这里说一个实用技巧训练前不要开启太多无关程序。我在实际项目中发现开着浏览器训练显存占用会多出几百兆对于显存比较小的显卡来说可能就是压垮骆驼的最后一根稻草。用nvidia-smi查看显存占用确保训练前GPU是干净状态。6.3 检测漏检和误检的处理模型训练完放到真实场景一测通常会发现两个问题漏检该检测出来的没检测出来和误检把不是头盔的当成头盔。漏检的处理有几个方向。第一降低置信度阈值模型输出更多候选框漏检自然少。第二检查是不是目标太小如果是远距离的目标检测不到考虑用更大尺寸训练或者提升图像的采集分辨率。第三补充更多遮挡、模糊、暗光环境下的训练数据这是提升模型泛化能力最根本的手段。误检的处理思路不同。误检通常是模型把和头盔外观相似的物体比如圆的铁桶、白色管道、安全帽形状的装饰物识别成了头盔。处理方法也是三个方向提高置信度阈值、补充难负样本把容易误检的图片加入训练集但标注为空告诉模型这些不是头盔、在业务逻辑上增加过滤条件。另外混淆矩阵能帮你定位到具体是哪两类在互相干扰针对性地加数据往往比盲目调参更有效。6.4 实测踩坑汇总整理几个实操中容易踩的坑每一个都是我付出过时间成本换来的。第一个坑验证集效果很好一到真实场景就翻车。原因基本是训练数据和真实场景分布不一致。比如训练数据全是白天、晴天、正面拍摄真实场景却是傍晚逆光、远距离俯拍。解决办法是数据采集时尽可能贴近真实部署环境模拟各种光照和角度。第二个坑用OpenCV读取RTSP视频流时经常断流或者延迟很大。这个不完全是模型的问题我后来把RTSP的传输协议从rtsp改成rtsptTCP传输稳定性有明显提升。还有就是在读取视频帧的循环里加一个缓存队列避免因为处理速度跟不上导致画面延迟累积。第三个坑多个摄像头同时检测时单卡GPU算力不够。我们的对接方案是GPU分时复用或者把推理任务拆到多块GPU上每块GPU负责一路摄像头的检测。如果设备端算力有限还可以降低每秒处理的帧数比如每两帧抽一帧检测对头盔检测这种安全场景来说每秒5帧左右的检测频率已经足够及时发现违规行为。第四个坑是数据泄露。这个比较隐蔽我在一次复盘中发现某个类别的验证集AP异常高排查后意识到训练集和验证集里出现了同一段视频的连续帧模型相当于记住了答案而不是学会了识别。后来我就坚持做去重绝不让验证集内容混入训练集。这个原则看起来简单实际操作中非常容易疏忽。最后再分享一个小技巧训练完不要急着删掉所有实验记录runs/train/下的每个exp目录都保留。后面做模型迭代和效果对比时有历史实验做参照你能很清楚地知道哪个改动带来了收益、哪个改动的效果是负面的。我做头盔检测这个项目前后产生了二十多个实验记录最后能快速锁定最佳模型靠的就是这套留痕习惯。头盔检测这个方向的技术难度不算高但真正落地时要考虑的东西不少从数据分布到推理性能从告警逻辑到设备适配每个环节都值得花时间打磨。这套基于YOLOv5和PyTorch的方案已经在我们几个项目里稳定跑了很久如果你也在做类似的安全检测场景可以参考这套流程直接复现然后在数据层面多下功夫效果应该不会让你失望。本文还有配套的精品资源点击获取
简介目标检测是计算机视觉的核心任务之一旨在从图像或视频中定位并识别特定对象。YOLOv5作为单阶段检测算法的代表通过回归方式直接输出目标位置与类别在保持高精度的同时实现了极快的推理速度尤其适用于需要实时响应的安全监管场景。PyTorch则凭借动态计算图和成熟的CUDA支持为模型训练与调试提供了极大便利。基于这两项技术构建的头盔检测系统可自动识别监控画面中人员是否佩戴安全帽一旦发现违规立即告警广泛应用于工地、园区、工厂等安全生产领域。本文从环境配置、数据标注、模型训练到推理部署完整呈现了构建此类系统的全过程并分享了解决漏检、误检及显存不足等实际问题的经验为开发者提供了一套可复现的工程实践方案。 去年接了一个工地的安全监管需求靠人盯监控看工人有没有戴安全帽眼睛根本看不过来。后来就把这个需求落地成了基于YOLOv5和PyTorch的深度学习头盔检测系统能做到对图片、视频、实时监控画面中的头盔佩戴情况进行自动识别检测到未佩戴立刻框出来并告警。整套流程从环境搭建、数据标注、模型训练到最终部署推理我完整跑了一遍这篇文章就把里面的关键细节和踩过的坑一次性讲清楚。适合想用深度学习做目标检测但还停留在看过教程、没跑过完整项目阶段的人也适合准备在工地、园区、工厂等场景落地安全帽检测的开发者参考。1. 项目整体设计与技术选型思路1.1 检测算法选型为什么锁定YOLOv5头盔检测本质上是一个目标检测任务核心要解决两个问题第一把画面里的人找出来第二判断这个人头部区域有没有佩戴头盔。市面上可选的目标检测算法不少两阶段的Faster R-CNN、单阶段的SSD和YOLO系列我最终选了YOLOv5原因很直接。YOLO系列发展到v5这代工程化程度已经非常高。相比Faster R-CNN那种先提候选区域、再逐个分类回归的两阶段方案YOLOv5把检测当成一个回归问题一次前向推理直接输出所有目标的位置和类别速度优势非常明显。对于头盔检测这种需要实时响应的场景推理速度是硬指标一秒钟处理几十帧画面才有实用价值。另外YOLOv5的代码结构在同期开源项目里属于最清晰的那一档。模型定义、训练流程、数据加载、推理脚本全部模块化改配置就能换模型尺寸和数据路径这对我们这种要把项目真正落地、后期还要根据不同摄像头场景反复调优的团队来说太重要了。你想换YOLOv5s还是YOLOv5m改一行参数就行不用动代码。还有一个关键因素预训练权重。YOLOv5官方提供了在COCO数据集上预训练好的权重虽然COCO里没有头盔这个类别但模型已经学到了通用的特征提取能力——边缘、纹理、形状这些底层特征我们只需要在它的基础上用自己的头盔数据集做微调训练时间和数据量需求都能大幅降低。这跟让一个见过几万种物体的人去学认一个新物种比让婴儿从零开始学要快得多是一个道理。1.2 深度学习框架选型PyTorch为什么顺手框架选了PyTorch基本没纠结。一个项目里框架的选择往往决定了开发效率的上限PyTorch最大的优势是动态计算图模型结构可以在运行时改变调试的时候你能随时打印每一层的输出形状和梯度信息遇到问题可以非常直观地定位。举个例子训练过程中如果某一层的输出维度对不上PyTorch会直接在报错信息里告诉你张量的形状和预期值配合print或者pdb断点几乎不需要猜。相比之下静态图框架是先构图、后执行调试时候的反馈链路会绕很多。对深度学习的初学者和中等水平开发者来说PyTorch的学习曲线也友好得多社区里能找到的海量教程和开源代码大部分都是PyTorch版本遇到问题一搜基本都有答案。还有一点是硬件生态。我们当时训练用的就是NVIDIA的GPUPyTorch对CUDA的支持非常成熟从显卡驱动到CUDA Toolkit再到PyTorch版本官方文档给出了完整的对应关系。训练时一个model.cuda()就能把网络搬到GPU上nvidia-smi可以实时查看显存占用非常顺手。为什么用PyTorch而不用TensorFlow这个问题我自己的答案很朴素YOLOv5官方源码本身就是基于PyTorch实现的选框架的第一个原则是跟随项目生态能直接跑通、能快速迭代比什么都重要。框架本身只是个工具好用、社区活跃、坑少就是最大的竞争力。1.3 系统整体架构与工作流程整个头盔检测系统可以拆成离线训练和在线推理两条链路。离线训练链路是数据采集和标注、数据集划分、模型训练、效果评估产出一个训练好的权重文件。在线推理链路是读取图片或视频流、送入模型做前向推理、后处理非极大值抑制、在画面上绘制检测框和置信度、根据检测结果触发告警逻辑。两条链路都围绕同一套模型文件运转这也是我比较推荐的落地方式。训练和推理解耦以后你可以在训练阶段随意折腾模型结构和超参数一旦权重固定下来推理端就可以做成一个轻量服务比如用Flask对模型做一层封装再对接监控系统的告警模块。我们当时的告警逻辑做了个简单的业务规则检测到人但没检测到头盔且人的检测框与头盔检测框没有重叠就判定为未佩戴触发记录和提醒。这里有个细节直接用人和头盔两个类别的检测结果做逻辑判断比单独训练一个佩戴/未佩戴分类模型要灵活得多因为画面里人很多逐个分类会增加漏检风险而目标检测可以同时输出所有人的位置和属性规则也好调整。2. 环境搭建从零配置PyTorch与YOLOv52.1 深度学习环境的三层结构环境搭建往往是最劝退新手的一步很多人在这一步就被各种版本冲突搞得心态崩溃。我把深度学习环境拆成三个层级来理解会清晰很多。最底层是驱动和CUDA。显卡驱动是操作系统和GPU硬件之间的桥梁CUDA Toolkit是NVIDIA提供的并行计算平台PyTorch在GPU上跑底层就是调用CUDA来操作GPU算力。往上走是Python环境我强烈建议用Anaconda来管理每个项目创建独立的虚拟环境互不干扰。最上层才是PyTorch框架本身。一个常见的误解是CUDA装得越新越好。其实不是关键是驱动版本要支持你选的CUDA版本而PyTorch版本又要对应匹配某个CUDA版本。这块我踩过一次坑有台机器装了很新的驱动我直接装了最新版CUDA结果那个版本的PyTorch还不支持又得重来一遍。后来我学乖了一律按PyTorch官方安装命令里指定的CUDA版本来驱动驱动的安装。用nvidia-smi查看显卡驱动版本右上角的CUDA Version表示该驱动最高支持的CUDA版本只要你的CUDA Toolkit版本不高于这个数字基本就没问题。然后去PyTorch官网选对应组合官网会给出安装命令直接复制执行即可。2.2 PyTorch GPU版本的安装与验证我们当时的安装路线是在Ubuntu 22.04上配置的命令大概是这样的conda create -n helmet python3.9 conda activate helmet pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118这里有两个点要特别说明。第一创建环境时Python版本选了3.9不是最新的3.11或者3.12原因是YOLOv5的依赖库尤其是numpy、opencv-python这些在3.9下兼容性最稳。深度学习项目的依赖往往很敏感Python小版本升级可能导致某个库没有对应的预编译包所以建议用版本适中、生态成熟的Python环境。第二PyTorch的安装命令里cu118表示CUDA 11.8版本。为什么选11.8而不是更高的12.x因为当时YOLOv5的官方测试环境就是11.8社区反馈这个组合下各种算子都稳定没必要追求最新。安装完成之后一定要验证GPU能不能正常调用我一般会跑这三行import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))如果输出True和显卡型号说明PyTorch已经能正常调用GPU。这一步验证不能省很多人装完就急着训练结果模型一直在CPU上跑速度慢得离谱还以为是代码问题。2.3 获取YOLOv5源码与安装依赖PyTorch环境就绪后接着获取YOLOv5源码。这里我建议直接用Git拉官方仓库后续要更新代码或者切换版本也方便git clone https://github.com/ultralytics/yolov5 cd yolov5 pip install -r requirements.txtrequirements.txt里包含了YOLOv5运行所需的全部依赖包括numpy、opencv-python、matplotlib、seaborn这些。有些教程会让你一条条手动安装没必要一条命令搞定。装依赖的时候有个小坑默认PyPI源在国内拉取很慢有些大包比如opencv-python动辄几十兆容易超时。我的做法是直接用国内镜像源把pip下载地址切到清华源或者阿里源速度能快一个数量级pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple环境装好后可以直接跑一下官方自带的推理脚本用预训练权重yolov5s.pt去检测自带的测试图片确认整个链路是通的再开始准备自己的数据。这一步相当于验收环境确保问题不堆积到后续环节。3. 头盔数据集准备与标注3.1 数据集来源与规模规划深度学习是数据喂出来的模型效果好不好的下限由数据质量决定。头盔检测的数据集有两个主要来源一是公开数据集二是自己采集标注。我们在项目初期先用公开数据集把整套训练流程跑通后期再针对实际场景补充自采数据。公开数据集里有几个可以重点关注SCUT-HEAD这是华南理工大学开源的头部检测数据集包含超过4000张标注好的头部图像还有其他研究者发布的建筑工地安全帽佩戴检测数据集标注了head和helmet两个类别。我当时的做法是把几个公开数据集合到一起做了一个大约8000张图的初始数据集类别就两类helmet和head。这里有个很关键的经验不要贪多求全。头盔检测的类别设计越简单越好我见过有人把类别分得很细什么红色安全帽白色安全帽蓝色安全帽结果每个类别的样本量都不够模型训练出来mAP惨不忍睹。头盔就是头盔检测目标就是有没有戴类别分多了只会增加模型的负担。数据集规模上头盔检测这类任务两张类别每类3000到5000个标注实例已经能训练出一个可用的模型。如果完全从零开始训练没有预训练权重可能需要更多数据但我们用了YOLOv5的COCO预训练权重做迁移学习数据需求量就小很多。3.2 数据标注labelImg标注头盔类别数据标注是项目里最枯燥但最重要的一环直接决定模型精度的上限。我们用的标注工具是labelImg一个基于Python的图形化标注工具体积小、操作简单、跨平台标注结果直接保存为YOLO格式的txt文件。安装labelImg很简单conda install pyqt5 pip install labelImg labelImg打开工具后用Open Dir打开图片文件夹然后通过Change Save Dir设置标注文件保存路径再在左侧选择类别最常用的快捷键W是创建标注框A和D切换上一张/下一张图CtrlS保存。标注的核心原则是框得准、不要漏。画框的时候框体要紧贴目标边缘稍微留一点点边距可以但不要框得太大把背景包进去也别框得太小截断了目标的一部分。一个容易忽略的细节是画面里非常小、非常模糊的目标也要尽量标注因为现实场景中摄像头距离远的时候目标就是很小的一块如果训练数据里全是清晰的大目标模型在真实场景下的召回率会很低。3.3 数据集目录组织与data.yaml配置标注完成后数据集要按YOLOv5要求的目录结构组织起来。YOLOv5的数据加载逻辑是images目录放图片labels目录放标注文件图片和标注文件通过文件名一一对应后缀不同.jpg和.txt名称相同。目录结构如下helmet_dataset/ ├── images/ │ ├── train/ # 训练集图片 │ └── val/ # 验证集图片 └── labels/ ├── train/ # train图片对应的标注txt └── val/ # val图片对应的标注txt训练集和验证集的划分我习惯按9:1或者8:2来但划分之前必须做一步去重。因为很多公开数据集的图片往往是从视频里逐帧截取的连续帧之间相似度极高如果不做去重模型会在验证集上表现很好到了真实场景就打回原形这就是过拟合的典型表现。我的做法是用脚本计算图片的感知哈希把相似度高的图片归到同一个集合里确保训练集和验证集的内容足够独立。标注文件的内容是YOLO格式的纯文本每一行代表一个目标对象格式为类别id x_center y_center width height四个坐标值都是归一化到0到1之间的比例值。比如0 0.5125 0.3452 0.2341 0.1876表示一个类别为0的目标中心点在图片的横向51.25%、纵向34.52%的位置宽高分别占整图的23.41%和18.76%。最后配置一个data.yaml文件内容很简单train: /path/to/helmet_dataset/images/train val: /path/to/helmet_dataset/images/val nc: 2 names: [helmet, head]nc是类别数names是类别名称列表顺序必须跟标注文件里的类别id一致。这个data.yaml在训练时用--data参数指定YOLOv5会自己读取。4. 模型训练参数理解与效果调优4.1 从卷积到检测YOLOv5核心结构速览训练之前哪怕不打算改模型结构也建议把YOLOv5的网络结构过一遍否则调参的时候会一头雾水。YOLOv5整体可以分成三块Backbone、Neck和Head。Backbone负责提取图像特征使用的是CSPDarknet结构。它的特点是引入了Cross Stage Partial结构把特征图分成两部分一部分走卷积层做特征提取另一部分直接连接到后面既减少了计算量又解决了深层网络梯度消失的问题。你可以把它理解为一条高速公路信息既能走铺满细节的小路也能走无障碍的直达通道。Neck部分用的是PANet结构它的作用是融合不同尺度的特征。头盔有大有小——近处的头盔占半个屏幕远处的头盔可能就十几个像素。PANet通过自顶向下和自底向上的双向路径把浅层的细节信息和深层的语义信息融合在一起让模型对小目标更敏感。Head部分输出三个不同尺寸的检测结果分别对应大、中、小目标。这也是YOLO系列的经典设计每个网格单元预测若干个锚框每个锚框输出位置偏移量、宽高、置信度和类别概率。至于深度学习的池化这个知识点在YOLOv5中也扮演了角色。SPPSpatial Pyramid Pooling模块在Backbone的最后一层使用不同池化核尺寸对特征图做池化再把结果拼接起来目的是增强模型对不同尺度目标的感受野。简单说池化就是压缩信息、扩大视野让网络从更大范围去理解一个目标是什么。4.2 训练超参数每一行配置背后的逻辑训练头盔检测模型最常用的是YOLOv5s这个尺寸的模型。yolov5s是官方提供的五种尺寸中最小的速度最快精度也不错。对头盔检测这种目标特征相对简单的任务用yolov5l甚至yolov5x属于杀鸡用牛刀训练时间长、推理速度慢精度提升却很有限。启动训练的命令cd yolov5 python train.py --data helmet.yaml --cfg yolov5s.yaml --weights yolov5s.pt --epochs 100 --batch-size 16 --img 640几个关键超参数我逐个说一下。--epochs是训练轮数。100轮是一个比较合理的起点YOLOv5训练过程中会自动保存效果最好的权重所以轮数多一点不用担心过拟合配合早停机制就行。我见过有人只训练30轮就停下来用效果好不到哪去因为模型还没充分收敛。--batch-size是批次大小。这个参数受限于显存显存不够就调小一点。但注意batch-size大小会直接影响训练效果太大的batch会让模型收敛不稳定太小的batch又会导致梯度估计噪声大。--img是输入图片的尺寸。默认640x640这个尺寸对大多数场景够用。头盔检测如果主要看监控画面目标通常比较小可以考虑用更大尺寸比如960或者1280来训练对小目标检测有提升但对应的显存消耗和推理耗时也会增加。还有一组藏在data/hyps/hyp.scratch-low.yaml里的超参数它们是深度学习训练中容易被忽略但很关键的细节。比如lr0是初始学习率momentum是动量weight_decay是权重衰减系数。这些参数YOLOv5官方已经调得比较合理新手不建议动等你跑完一轮训练对损失下降曲线有感觉了再来调学习率也不迟。4.3 启动训练与过程监控训练跑起来后不是干等着要会看训练日志。YOLOv5每个epoch结束会输出一组指标包括box_loss、obj_loss、cls_loss和mAP。box_loss代表预测框位置和真实框的误差这个值应该在训练中逐渐下降obj_loss代表目标置信度误差也需要下降cls_loss是分类误差对头盔检测来说前期会降得比较快因为两类区分度很高。mAP0.5是最直观的指标表示IoU阈值0.5下的平均精度均值头盔检测任务一般能训练到0.85以上就算非常可用的模型了。训练过程中有一个现象第一次遇到可能会慌损失值在初始几个epoch会跳来跳去甚至短暂上升这是正常的。因为迁移学习初期模型在快速调整底层特征损失波动不代表训练失败。我的经验是前20轮不用频繁关注指标到30轮以后再看趋势才有意义。显存不足的情况在训练中很常见报错通常是CUDA out of memory。最简单的处理是调小--batch-size从16降到8或者4如果还是不足就换小尺寸的模型。另外可以在训练命令中加上--cache参数把数据预加载到内存中能减少GPU和CPU之间的数据搬运等待但并不降低显存占用。4.4 训练结果评估mAP、PR曲线怎么读训练结束后YOLOv5会在runs/train/目录下生成实验记录里面有results.png、confusion_matrix.png、PR_curve.png等图表这些图对一个深度学习项目来说就是体检报告。PR_curve.png是精度-召回率曲线。精度Precision衡量的是模型预测为头盔的框里有多少真正的头盔召回率Recall衡量的是真实头盔中有多少被模型找出来了。这两个指标天然矛盾鱼和熊掌不可兼得。PR曲线与坐标轴围成的面积就是AP值面积越大表示模型效果越好。头盔检测这种安全相关的场景我个人更看重召回率——漏检一个没戴头盔的人比多误报几次更严重。所以后期调参时可以适当降低置信度阈值提高召回率。confusion_matrix.png是混淆矩阵能直观看出模型在哪两类之间混淆。对头盔检测来说最常见的混淆就是头盔和头部互相误判比如把戴浅色帽子的人的头误判为头盔。如果混淆矩阵里这两类的交叉区域比较大说明训练数据里头盔和头部的外观特征区分度不够需要补充更多难例数据。还有一个容易被忽视的评估指标是各类别的AP单独值。YOLOv5训练日志里会按类别分别输出AP如果helmet类AP很高但head类AP很低说明数据集中head类的样本不够或者标注质量有问题需要针对性补充。5. 实时检测从模型到应用5.1 使用训练好的权重进行图片和视频推理训练完成后最佳权重保存在runs/train/exp/weights/best.pt。使用这个权重做推理很简单python detect.py --weights runs/train/exp/weights/best.pt --source data/images --conf-thres 0.5--source既可以是图片文件夹也可以是一个视频文件路径甚至是摄像头设备的索引号比如--source 0表示读取本机第一个摄像头。推理时的一个关键调整参数是--conf-thres即置信度阈值。这个值的设置直接影响系统的误报和漏报阈值设高比如0.7模型只输出有把握的结果误报少但容易漏检阈值设低比如0.3漏检少但误报多。我建议头盔检测场景设0.4到0.5之间兼顾两种情况。detect.py跑完后结果图片和视频会保存在runs/detect/目录下检测框会标注类别名称和置信度。这里有个小技巧可以给不同类别设置不同颜色和线宽让未佩戴头盔的人脸框显示为红色醒目提示而已佩戴的框显示为绿色这对接实际业务体验很重要。5.2 实时视频流检测与部署detect.py默认方式是对视频逐帧检测然后保存结果但真正的实时检测系统需要把模型集成到自己的代码流程里。核心思路是先加载模型然后循环读取视频帧每帧做一次前向推理。一个简化的检测循环框架import torch import cv2 model torch.hub.load(ultralytics/yolov5, custom, pathruns/train/exp/weights/best.pt, force_reloadTrue) model.conf 0.5 model.iou 0.45 cap cv2.VideoCapture(0) # 读摄像头或视频流 while True: ret, frame cap.read() if not ret: break results model(frame) # results.pandas().xyxy[0] 包含所有检测框的坐标、类别和置信度 # 在这里编写告警逻辑有人但无头盔 - 发送告警 rendered results.render()[0] cv2.imshow(Helmet Detection, rendered) if cv2.waitKey(1) 0xFF ord(q): breaktorch.hub.load是PyTorch提供的一个非常方便的功能可以从本地加载自定义训练好的模型不需要手动构造网络结构。模型加载后通过设置model.conf和model.iou控制置信度阈值和NMS的IoU阈值这也是在线部署时最常用到的两个后处理参数。实时检测的实测性能需要心里有数。YOLOv5s在NVIDIA GTX 1660 Ti显卡上640x640输入推理速度大概在3到5毫秒一帧也就是每秒可以处理200到300帧实时性完全不是瓶颈。真正的瓶颈往往在视频解码和显示环节cv2.VideoCapture读RTSP流时延迟可能达到几百毫秒这是摄像头侧和网络侧的问题跟模型无关。5.3 模型导出与边缘设备部署模型训练好之后如果只在自己电脑上跑演示没问题但实际项目大多要部署到边缘设备或者服务器上。YOLOv5官方支持将模型导出为多种格式ONNX、TensorRT、CoreML、OpenVINO。导出ONNX格式的命令python export.py --weights runs/train/exp/weights/best.pt --include onnx --img 640导出ONNX的意义在于跨平台部署。ONNX是一种开放的模型表示格式相当于深度学习的通用语言很多推理引擎和边缘设备芯片比如RK3568、Jetson系列都支持加载ONNX模型。我们在项目中用了瑞芯微的RK3568平台做边缘部署流程是PyTorch训练 - 导出ONNX - 转换成RKNN格式 - 在板子上跑推理。导出时有个细节需要留意--img 640要和训练时的输入尺寸一致否则导出后模型可能因为输入尺寸不匹配而报错或者精度下降。另外如果想在TensorRT上部署建议先导出ONNX再转TensorRT引擎而不是直接导出因为兼容性更好。Jetson设备上部署PyTorch模型是另一个常见需求。注意Jetson平台使用的JetPack版本决定了支持的PyTorch版本比如JetPack 6.x一般对应PyTorch 2.x安装时不要用桌面版通用的pip命令而是参考NVIDIA官方针对Jetson提供的预编译wheel包这样能避免很多兼容性问题。6. 常见问题与排查技巧实录6.1 训练不收敛的排查思路训练过程中最让人崩溃的就是loss不降或者降着降着突然变成NaN非数值然后训练崩了。这个问题我遇到过不止一次原因是多方面的排查顺序建议按照数据 - 超参数 - 环境来。先说数据层面。检查标注文件里有没有异常的框比如坐标值超过1、宽高为0这种脏数据。YOLO格式要求归一化坐标在0到1之间如果标注工具或脚本出了问题某个坐标写成了5.7模型训练时梯度就会爆炸。我有一个脚本专门扫描数据集里的非法标注跑一遍就能发现并剔除问题样本。其次是学习率。如果loss前几轮就出现NaN大概率是学习率太大。YOLOv5的默认初始学习率是0.01如果你的数据集很小可以把hyp.scratch-low.yaml里的lr0调到0.005再试。学习率调小了一般能解决代价是收敛速度变慢多跑几轮就行。环境层面的问题比较隐蔽。有一种情况是GPU驱动和CUDA版本不匹配导致某些算子计算异常表面上看loss一切正常但训练到中途就开始跳NaN。检查方法是先用官方模型跑一个小的训练任务比如只训练5轮、只用100张图如果官方模型也崩那问题就不在你自己的数据和代码而是环境层面回退CUDA或者PyTorch版本重试。6.2 显存不足的解决路径CUDA out of memory是新手最常遇到的报错看到一堆红色报错信息容易慌其实解决路径很清晰。第一步把batch-size调小。训练命令中--batch-size 16改成8如果还爆就改成4。batch-size是显存占用的主要因素调小它效果立竿见影。第二步是换小模型YOLOv5s已经是小模型了如果用的yolov5m或者更大换回yolov5s。第三步是减小输入图片尺寸--img 640改成--img 416显存占用会下降不少但检测精度也会相应降低。这里说一个实用技巧训练前不要开启太多无关程序。我在实际项目中发现开着浏览器训练显存占用会多出几百兆对于显存比较小的显卡来说可能就是压垮骆驼的最后一根稻草。用nvidia-smi查看显存占用确保训练前GPU是干净状态。6.3 检测漏检和误检的处理模型训练完放到真实场景一测通常会发现两个问题漏检该检测出来的没检测出来和误检把不是头盔的当成头盔。漏检的处理有几个方向。第一降低置信度阈值模型输出更多候选框漏检自然少。第二检查是不是目标太小如果是远距离的目标检测不到考虑用更大尺寸训练或者提升图像的采集分辨率。第三补充更多遮挡、模糊、暗光环境下的训练数据这是提升模型泛化能力最根本的手段。误检的处理思路不同。误检通常是模型把和头盔外观相似的物体比如圆的铁桶、白色管道、安全帽形状的装饰物识别成了头盔。处理方法也是三个方向提高置信度阈值、补充难负样本把容易误检的图片加入训练集但标注为空告诉模型这些不是头盔、在业务逻辑上增加过滤条件。另外混淆矩阵能帮你定位到具体是哪两类在互相干扰针对性地加数据往往比盲目调参更有效。6.4 实测踩坑汇总整理几个实操中容易踩的坑每一个都是我付出过时间成本换来的。第一个坑验证集效果很好一到真实场景就翻车。原因基本是训练数据和真实场景分布不一致。比如训练数据全是白天、晴天、正面拍摄真实场景却是傍晚逆光、远距离俯拍。解决办法是数据采集时尽可能贴近真实部署环境模拟各种光照和角度。第二个坑用OpenCV读取RTSP视频流时经常断流或者延迟很大。这个不完全是模型的问题我后来把RTSP的传输协议从rtsp改成rtsptTCP传输稳定性有明显提升。还有就是在读取视频帧的循环里加一个缓存队列避免因为处理速度跟不上导致画面延迟累积。第三个坑多个摄像头同时检测时单卡GPU算力不够。我们的对接方案是GPU分时复用或者把推理任务拆到多块GPU上每块GPU负责一路摄像头的检测。如果设备端算力有限还可以降低每秒处理的帧数比如每两帧抽一帧检测对头盔检测这种安全场景来说每秒5帧左右的检测频率已经足够及时发现违规行为。第四个坑是数据泄露。这个比较隐蔽我在一次复盘中发现某个类别的验证集AP异常高排查后意识到训练集和验证集里出现了同一段视频的连续帧模型相当于记住了答案而不是学会了识别。后来我就坚持做去重绝不让验证集内容混入训练集。这个原则看起来简单实际操作中非常容易疏忽。最后再分享一个小技巧训练完不要急着删掉所有实验记录runs/train/下的每个exp目录都保留。后面做模型迭代和效果对比时有历史实验做参照你能很清楚地知道哪个改动带来了收益、哪个改动的效果是负面的。我做头盔检测这个项目前后产生了二十多个实验记录最后能快速锁定最佳模型靠的就是这套留痕习惯。头盔检测这个方向的技术难度不算高但真正落地时要考虑的东西不少从数据分布到推理性能从告警逻辑到设备适配每个环节都值得花时间打磨。这套基于YOLOv5和PyTorch的方案已经在我们几个项目里稳定跑了很久如果你也在做类似的安全检测场景可以参考这套流程直接复现然后在数据层面多下功夫效果应该不会让你失望。本文还有配套的精品资源点击获取
平方复杂度无法被常规优化消除｜TaoToken 统一 Key 通道实测)
