基于改进YOLOv5s和TensorRT部署的鱼道过鱼监测

作者：李健源; 柳春娜; 卢晓春; 吴必朗

来源：农业机械学报, 2022, 53(12): 314-322.

DOI：10.6041/j.issn.1000-1298.2022.12.031

• 鱼道过鱼监测
• YOLOv5s
• 损失函数
• TensorRT

为实现在复杂水体下对鱼道过鱼进行监测，提出了一种基于YOLOv5s的改进模型，并用TensorRT部署应用于某水电站鱼道现场。首先，针对水下图像模糊、目标检测困难的问题，提出了将Swin Transformer(STR)作为检测层，提升了模型对目标的检测能力；其次，针对鱼群密集、图像信息少的问题，将Efficient channel attention(ECA)注意力机制作为主干特征提取网络C3结构的Bottleneck,减少了计算参数并提升了检测精度；然后，针对检测目标定位差、回归不收敛的问题，将Focal and efficient IOU loss(FIOU)作为模型损失函数，优化了模型整体性能；最后将模型部署在TensorRT框架进行优化，处理速度得到了大幅度提升。基于实际采集的复杂水体数据集进行实验，结果表明，本文算法mAP为91.9%,单幅图像处理时间为10.4 ms,在相同条件下，精度比YOLOv5s提升4.8个百分点，处理时间减少0.4 ms。模型使用TensorRT部署后单幅图像推理时间达到2.3 ms,在不影响检测精度的前提下，推理速度提高4.5倍。综上，本文算法模型在保证快速检测的基础上，具有较高的准确性，更适用于复杂水体下鱼道过鱼监测。

• 单位
  三峡大学; 中国水利水电科学研究院; 流域水循环模拟与调控国家重点实验室