Yolov8 目标检测蒸馏学习记录
yolov8系列模型蒸馏基本流程,代码下载:这里本人提交了一个demo:djdll/Yolov8_Distillation: Yolov8轻量化_蒸馏代码实现
在轻量化模型设计中,**知识蒸馏(Knowledge Distillation)**被广泛应用,作为提升模型性能的有效手段。它通过引入一个性能更强的教师模型,指导学生模型在保持高推理速度和小模型规模的前提下提升精度。在目标检测任务中,蒸馏技术特别适用于精度要求高但部署资源受限的场景,如边缘设备、工业检测等。检测主流的包含输出蒸馏、和特征蒸馏以及现在还有两者结合的方法。
输出蒸馏也被称为logits蒸馏或head-level蒸馏
它直接作用于检测器的最终输出层,指导学生学习教师的预测结果。对于目标检测模型来说,这些输出主要包括:
- 类别概率分布(cls)
- 目标置信度(obj/conf)
- 边界框位置(bbox)
通过最小化学生输出与教师输出之间的差异,学生能够从教师的“软标签”中学习更加细致的类别区分和更准确的目标定位信息。这种方法实现简单,对结构差异容忍度高,是目标检测蒸馏中应用最广泛的形式之一。
另一种特征蒸馏(Feature-Based Distillation):
特征蒸馏的核心思想是让学生模型在中间层提取到与教师模型相似的语义表示,从而提升整体特征表达能力。常见做法包括:
- 单层对齐:选择某一关键层(如 Backbone 的最后一层)作为对齐目标;
- 多层对齐:对多个阶段或模块(如 FPN、Neck)同时进行蒸馏,以增强全局信息传递;
- 注意力蒸馏:通过显式提取教师的注意力图(通道或空间注意力)作为引导,强化学生对关键区域的感知能力。
特征蒸馏一般需要对齐教师和学生的中间特征维度,可能引入额外的投影模块(如 1×1 卷积)进行通道匹配。
实践示例:以 YOLOv8 检测为例
在 YOLOv8 中,蒸馏的实现流程通常包括以下几个步骤:
- 构建教师模型与学生模型,分别加载预训练权重。
- 前向传播中同时计算教师和学生的输出,提取检测头与特征图。
- 设计蒸馏损失函数
- 将蒸馏损失与原始检测损失加权融合,参与总 loss 回传更新学生模型。
- 仅保存训练好的学生模型用于部署。
蒸馏训练
- 以yolov8s和yolov8n分别作为教师和学生模型训练代码
import os
from ultralytics import YOLO
import torchos.environ['KMP_DUPLICATE_LIB_OK'] = 'True'def main():model_t = YOLO('runs/detect/train_v8s/weights/best.pt') # the teacher modelmodel_s = YOLO('runs/detect/train_v8n/weights/best.pt') # the student model"""Attributes:Distillation: the distillation modelloss_type: mgd, cwdamp: Automatic Mixed Precision"""model_s.train(data="ultralytics/cfg/datasets/coco128.yaml", Distillation=model_t.model, loss_type='mgd', amp=False, imgsz=640, epochs=100,batch=32, device=0, workers=0, lr0=0.001)if __name__ == '__main__':main()
- 如果教师和学生模型并非v8s和v8n需要调试更改下面代码
在ultralytics/engine/trainer.py中FeatureLoss类下的def forward(self, y_s, y_t):处进行断点调试,分别得到y_s和y_t对应的通道数
然后将对应的通道写入Distillation_loss类
channels_s = [128, 256, 128, 64, 128, 256][-le:]
channels_t = [256, 512, 256, 128, 256, 512][-le:]
最后重新训练模型即可
结果
- 以coco128训练,为了快速验证,教师模型和学生模型各训练50e,btach 32,蒸馏训练50e,batch 32。
- yolov8s教师模型
- yolov8n学生模型
- 蒸馏模型
实验记录,不做参考,主要目的是跑通流程,具体还要在自己数据集上实测。
参考
本实现参考开源项目:https://github.com/jasonDasuantou/yolov8_distillation
蒸馏方法:MGD