YOLO-S:一种新型轻量的安全帽佩戴检测模型

摘要

针对现有施工场所下工人安全帽佩戴检测模型推理耗时长、对硬件要求高,且复杂多变环境下的训练数据集单一、数量少导致模型鲁棒性较差等问题,提出了一种轻量化的安全帽佩戴检测模型YOLO-S.首先,针对数据集类别不平衡问题,设计混合场景数据增强方法,使类别均衡化,提高模型在复杂施工环境下的鲁棒性;将原始YOLOv5s主干网络更改为MobileNetV2,降低了网络计算复杂度.其次,对模型进行压缩,通过在BN层引入缩放因子进行稀疏化训练,判定各通道重要性,对冗余通道剪枝,进一步减少模型推理计算量,提高模型检测速度.最后,通过知识蒸馏辅助模型进行微调得到YOLO-S.实验结果表明,YOLO-S的召回率及mAP较YOLOv5s分别提高1.9%、1.4%,模型参数量压缩为YOLOv5s的1/3,模型体积压缩为YOLOv5s的1/4, FLOPs为YOLOv5s的1/3,推理速度快于其他模型,可移植性较高.

关键词

安全帽佩戴检测 / 数据增强 / 模型压缩 / 知识蒸馏