基于强大的目标检测能力和较快的推理速度,YOLOv5 已经逐渐成为实时性目标检测任务中的首选模型。因此,YOLOv5 模型的优化和部署是落地任务中非常重要的部分。在 CPU 中部署 AI模型存在着大量的需求,Intel CPU 因其强劲的性能,丰富的软件生态,是我们在 CPU 上部署的首选硬件。本文主要研究了 YOLOv5 模型的知识蒸馏以及在 Intel CPU 上使用 OpenVINO 进行部署优化实践。通过采用知识蒸馏和量化的方法,在不损失精度的条件下,实现了 YOLOv5 模型在 OpenVINO 上 2 倍的推理速度提升。相关代码已开源在Adlik代码库中:https://github.com/Adlik/yolov5
1、知识蒸馏
现有的目标检测蒸馏方法主要对两阶段目标检测(如 RCNN 系列)较为有效,而这些蒸馏方法对于单阶段的目标检测器如 YOLOv5 模型精度几乎没有提升。因此,我们主要采用论文 《Object detection at 200 Frames Per Second》中对单阶段目标检测器的蒸馏方法。该论文研究了单阶段目标检测器使用普通知识蒸馏的问题和面临的挑战。
1.1 Objectness scaled Distillation
单级目标检测方法使用普通知识蒸馏的问题。单阶段目标检测器的预测是一个密集的候选集合。老师网络(YOLO 模型)预测图像背景区域中的边界框。在推理过程中,背景区域预测的边界框会被忽略。然而,标准的蒸馏方法会将这些背景检测转移到了学生模型学习中。它会影响边界框的训练 ,因为学生网络会从老师网络预测的背景区域中学习错误的边界框。两阶段的目标检测方法(如 RCNN)通过使用 RPN 网络来规避这个问题,因为 RPN 网络预测相对较少的候选区域。为了避免学习老师网络对背景区域的预测,论文中定义蒸馏损失为 objectness scaled function。其思想学生网络只学习老师网络预测的目标概率值较高的边界框位置和类别概率。
YOLO 目标损失函数由三部分组成:regression loss(回归损失函数), objectness loss (目标损失函数)和 classification loss(分类损失函数)。
目标损失函数部分不需要目标缩放,因为对于背景区域的预测的目标概率值比较低,因此,知识蒸馏的目标损失函数定义如下:
学生网络的目标缩放分类函数:
其中,上述损失函数的第一部分对应原始的目标分类损失函数,而第二部分是目标分类蒸馏损失部分。学生网络的目标缩放边界框位置损失函数:
老师网络将非常低的目标概率值作为大多数与背景相对应的预测框的权重。基于目标函数的缩放算法是一种用于单阶段目标检测器中蒸馏过滤器的缩放算法,因为该算法分配给背景预测框的权值很低。损失函数形式保持不变,但对于知识蒸馏损失函数,只是用老师网络的输出代替标签值。训练阶段总的损失函数为: