[论文解读] Knowledge Distillation in YOLOX-ViT for Side-Scan Sonar Object Detection

Martin Aubard, László Antal+2 more|arXiv (Cornell University)|Jan 1, 2024

Advanced Neural Network Applications被引用 5

一句话总结

本论文提出 YOLOX-ViT，一种基于视觉Transformer的侧扫声呐图像目标检测器，并采用知识蒸馏技术将其压缩为更小、更高效的模型。该方法将墙体检测的误报率最高降低20.35%，同时将mAP50提升0.13–0.16，表明ViT融合与知识蒸馏相结合可显著提升水下目标检测的准确率与鲁棒性。

ABSTRACT

In this paper we present YOLOX-ViT, a novel object detection model, and investigate the efficacy of knowledge distillation for model size reduction without sacrificing performance. Focused on underwater robotics, our research addresses key questions about the viability of smaller models and the impact of the visual transformer layer in YOLOX. Furthermore, we introduce a new side-scan sonar image dataset, and use it to evaluate our object detector's performance. Results show that knowledge distillation effectively reduces false positives in wall detection. Additionally, the introduced visual transformer layer significantly improves object detection accuracy in the underwater environment. The source code of the knowledge distillation in the YOLOX-ViT is at https://github.com/remaro-network/KD-YOLOX-ViT.

研究动机与目标

• 通过引入视觉Transformer增强的YOLOX架构，提升侧扫声呐图像中的目标检测性能。
• 探究知识蒸馏在不损失检测准确率的前提下减小模型尺寸的有效性。
• 提出一个专注于水下机器人墙体检测的新侧扫声呐数据集。
• 评估视觉Transformer层在低能见度水下环境中对特征提取与检测鲁棒性的影响。
• 通过大模型教师网络蒸馏，降低紧凑模型的误报率。

提出的方法

• 在主干网络的SPPBottleneck之后插入视觉Transformer（ViT）层，替代标准卷积块，以增强YOLOX架构。
• 使用标准交叉熵损失和IoU-based损失函数，训练更大的YOLOX-L和YOLOX-L-ViT模型作为'教师'模型。
• 采用知识蒸馏，损失函数为 ℒ = λ·ℒhard + (1−λ)·ℒsoft，其中ℒhard为真实标签的交叉熵损失，ℒsoft为来自教师模型的软标签蒸馏损失。
• 将YOLOX-Nano和YOLOX-Nano-ViT模型通过蒸馏压缩为更小的'学生'模型，重点聚焦于特征图与logit蒸馏。
• 仅在教师模型训练阶段使用数据增强，而学生模型则在有无在线数据增强的情况下进行训练。
• 采用基于提示的蒸馏策略，聚焦于主干网络中间特征图，以对齐学生与教师模型的表示。

实验结果

研究问题

• RQ1 知识蒸馏能否有效降低紧凑YOLOX模型在侧扫声呐目标检测中的误报率？
• RQ2 在YOLOX中集成视觉Transformer层是否能提升水下环境中的检测准确率与鲁棒性？
• RQ3 ViT与知识蒸馏的结合如何影响小型模型的mAP50与推理效率？
• RQ4 在线数据增强对本研究低数据场景下大模型与小模型性能的影响如何？
• RQ5 ViT层能否增强声呐图像中的特征提取能力，特别是对墙体等结构特征的检测？

主要发现

• 使用L-ViT教师模型时，知识蒸馏使YOLOX-Nano-noAug学生模型的误报检测数量减少了6%；而使用ViT增强的教师模型时，误报率降低了20.35%。
• YOLOX-Nano-ViT-noAug学生模型在视频中实现了38%的检测率，仅含13.9个误报，其精度与误报控制能力均优于基础Nano模型。
• YOLOX-L-ViT在无在线数据增强训练下，实现了0.41的mAP50与98.93%的检测时长，显著优于YOLOX-L（0.33 mAP50，87.05%检测时长）。
• 与非ViT对应模型相比，ViT层使L模型的检测准确率提升8%，Nano模型提升23%，尤其在长时视频推理中表现更优。
• 未使用在线数据增强训练的模型表现出更好的泛化能力与更低的过拟合现象，尤其在L与L-ViT变体中更为明显。
• ViT增强的学生模型（Nano-ViT-noAug）实现了0.13的mAP50与30.16%的精度，表明ViT即使在小型模型中也能提升特征学习能力。