☰
YOLOv5知识蒸馏实战:大模型当老师,小模型提精度,边缘部署提分2~4个点
2026/10/14 13:16:01 网站建设 项目流程

简介:一份面向目标检测开发者与模型轻量化研究者的YOLOv5知识蒸馏实战资源。常用于解决模型体积大、推理速度慢的落地难题,通过教师网络向学生网络迁移知识,压缩模型同时尽量保持检测精度。包内共8个文件,以4个zip压缩包、3个pt权重文件、1个py脚本为主,整体约593MB。zip中涵盖YOLOv5蒸馏工程代码、VOC格式数据及说明文档;pt权重对应教师与学生模型;py脚本负责数据预处理,便于直接复现实验。已有2024人学习下载。资源从蒸馏原理、环境配置到数据处理与完整训练代码均有覆盖,参照教程操作即可理解并掌握整套蒸馏流程,适用于想快速上手知识蒸馏并完成模型压缩的CV开发者和算法工程师。

1. 基于yolov5的知识蒸馏实战,先把大模型的底子搬给小模型

基于yolov5的知识蒸馏,简单讲就是拿一个大权重当老师、一个小权重当学生,让小的模型在推理速度和参数量不变的前提下,精度尽量向大模型靠近。比如 yolov5l 蒸馏 yolov5s,学生模型的部署资源占用几乎不变,mAP 却能比直接训练 yolov5s 高出 2~4 个点,这在边缘设备和低算力盒子上非常值钱。适合的人群是那些已经训练过 YOLO、能看懂 loss 和权重文件,但手头算力有限、正为「模型太大放不下、太小精度不够」发愁的从业者。这篇文章不是给你复述论文,而是给你一套能直接落地、把蒸馏 loss 做成独立模块的实战源码路线:从原理选型、最小跑通命令、核心代码,到 5 个高频翻车点逐一排查。

2. 检测任务的蒸馏为什么比分类难:先搞清蒸馏对象再选型

2.1 知识蒸馏的本质:温度、软标签和检测头的特殊结构

Hinton 提出的知识蒸馏,网上讲烂了的是那一套:老师网络的 logits 除以温度 T 再做 softmax,得到软标签,学生对软标签做 KL 散度,最后用 T 的平方补偿梯度幅度。这套思路在分类网络里直接套用毫无问题,因为输出就是一个 (B, num_classes) 的向量,蒸馏的对象非常干净。

到了 yolov5 这里就变了。检测头的输出不是单一 logits,而是三个尺度(P3/P4/P5)的特征图,每个位置的通道是 4 + 1 + C,分别对应回归坐标、objectness 置信度和类别概率。也就是说,同一个输出张量里混合了回归分支和分类分支,而这两个分支的概率分布含义完全不同。分类分支可以做 softmax 软标签蒸馏,回归分支却没有「概率」这一说,框坐标是连续值,没法直接拿 KL 散度去蒸。

所以检测任务里做蒸馏,第一件事就是拆:把输出分开,分类分支用 soft label 引导,回归分支用 L1 或 Smooth L1 只在前景位置计算。温度 T 的选择也要重新考虑。分类任务里常见的 T=20 在检测任务里基本不可用,因为检测头每个空间位置都有大量背景样本,背景类别占绝对主导,T 一高,所有分布都被拉平,学生只会学到「什么都差不多」,梯度信号被背景淹没。

我一般的习惯是:分类蒸馏 T 取 6~10,回归蒸馏不走 soft label 而是直接算回归值差异,这样最稳。

2.2 三种主流蒸馏范式:logits蒸馏、特征图蒸馏、解耦蒸馏

做知识蒸馏,你得先选定蒸馏什么。目前检测任务里常见的方案大致分三类,区别在「蒸馏对象」上。

第一类是 response-based,也就是直接蒸检测头的输出 logits。实现最简单,分类分支做 KL 散度,回归分支加一个 L1 约束。但问题在于,回归分支的 L1 约束其实没什么信息量,因为两个网络的框预测分布差异很大,强行对齐反而可能让学生学歪。第二类是 feature-based,把学生网络中间层输出的特征图对齐到老师的特征图,让学生的特征表达更接近老师。第三类是 decoupled knowledge distillation(解耦知识蒸馏),它把分类和回归彻底拆开,并且只在前景区域做回归蒸馏,分类蒸时用前景 mask 加权,强调困难样本和真实目标位置附近的信息。

我的选型结论很直接:yolov5 上做蒸馏,只做 logits 蒸馏收益很有限,feature-based 单独用又容易引入特征对齐的噪声,实战中最稳的组合是「decoupled 蒸馏 + 特征图对齐」双路并行,后面核心实现章节我会把这两块都写成代码。

2.3 teacher 和 student 的选配思路:不是越大越好

很多人一上来就喜欢用最大的模型当老师,用最小的模型当学生,觉得这样蒸馏效果一定最好。实际经验是:teacher 和 student 的能力差距太大,蒸馏反而会翻车。你让 yolov5x 去蒸 yolov5n,学生模型的容量根本装不下老师学到的全部知识,特征对齐时会强迫学生去拟合它做不到的分布,最后的结果往往比不用蒸馏还差。

我一般建议的原则是:teacher 与 student 的精度 gap 控制在 4~8 个点以内。比如 yolov5l 蒸 yolov5s 是最经典也最稳的组合,差距足够大、有蒸馏空间,但又不至于大到学生吸收不了。teacher 本身不能是过拟合严重的模型,否则它输出的软标签里带着大量错误信息,学生学到的就是老师的幻觉。建议用带 weight decay、训练充足的 teacher,如果在验证集上 teacher 自己都明显过拟合,先解决 teacher 的问题再谈蒸馏。选好了模型组合,下一步就是最小流程跑通。

3. 先把最小流程跑通:目录组织、teacher冻结和启动命令

3.1 项目目录与文件组织:把蒸馏源码独立成模块

拿到一套蒸馏实战源码,别急着改 yolov5 的训练主文件。我的习惯是把蒸馏能力做成独立的模块,挂在原有训练流程外面,而不是把 loss 逻辑混进 detect 层内部。这样好处很明显:基础训练代码几乎不动,回归到纯 yolov5 时也方便,蒸馏模块可以单独测试和调整。

一个我常用的目录结构如下:

distill_yolov5/ ├── models/ # yolov5 核心代码 ├── data/ # 数据集配置 yaml ├── weights/ │ ├── teacher.pt # 大模型训练好的权重 │ └── student.pt # 学生模型权重(初值或预训练) ├── losses/ │ ├── __init__.py │ ├── kd_loss.py # KL散度蒸馏 loss │ └── feature_loss.py # 特征图对齐 loss ├── utils/ │ └── torch_utils.py ├── train_distill.py # 蒸馏训练入口 └── distill.yaml # 蒸馏参数配置

train_distill.py 是整个流程的入口。它做的事情可以拆成四步:加载 teacher 并冻结、加载 student、前向时同时取两个模型的检测头特征、把蒸馏 loss 加到原始的 box_loss / cls_loss / dfl_loss 上。这样组织源码,后面调参和换模型都只用改 distill.yaml 和损失模块,不用在大文件里到处翻代码。

3.2 Teacher 权重加载与冻结:eval 模式 + 梯度关闭

蒸馏训练里 teacher 是只读的,它只负责前向输出指导和特征参考,不参与梯度更新。这一步写错会造成两后果:显存翻倍、BN 统计量被污染。

yolov5 权重加载用 attempt_load 即可,加载后立刻切 eval 模式,并把所有参数 requires_grad 置为 False。代码如下:

import torch from models.experimental import attempt_load def load_teacher(teacher_path, device): teacher = attempt_load(teacher_path, map_location=device) teacher.eval() for p in teacher.parameters(): p.requires_grad_(False) # 关键:如果 teacher 的 BatchNorm 还在训练模式, # 每个 batch 都会用 student 的统计数据污染 running_mean,必须关掉 return teacher

注意泛型不能用泛型;这段代码里最容易忽略的是teacher.eval()。如果你在训练循环里忘了切 eval,teacher 的 BN 层仍会更新 running 统计量,而且 teacher 前向的中间激活会保存下来用于反向传播。虽然 requires_grad 已经关闭了,但 PyTorch 仍可能缓存部分中间结果,显存占用会明显偏高。eval 模式下 teacher 的特征缓存可控,显存占用也稳定。

3.3 最小启动命令:先验证管道再谈效果

整个流程是否跑通,看两件事:loss 能正常下降、蒸馏 loss 的数值量级和基础 loss 在同一个量级。如果蒸馏 loss 直接压过基础 loss 几个数量级,说明权重参数没配好。先用一个小模型快速跑通,命令如下:

python train_distill.py \ --data data/voc.yaml \ --teacher weights/yolov5l.pt \ --student models/yolov5s.yaml \ --epochs 3 \ --batch-size 8 \ --kd-alpha 0.5 \ --kd-temperature 10.0 \ --feat-lambda 0.1 \ --device 0

这里的--kd-alpha 0.5表示蒸馏 loss 和基础 loss 按 0.5:1 的比例融合,--feat-lambda 0.1是特征对齐 loss 的权重,通常要比 KL 蒸馏小一个量级。--kd-temperature是分类分支蒸馏温度,10 是较安全的起点。第一次跑不建议直接把 epochs 拉到上百,用 3 个 epoch 验证数据加载、teacher 前向、student 前向三个管道是否正常。等到 loss 曲线稳定下降,再把训练完整跑起来。

4. 核心实现:特征对齐、KL蒸馏loss与挂接入训练循环

4.1 特征对齐模块:通道不同、尺寸不同怎么办

yolov5s 和 yolov5l 的特征图尺寸和通道数完全不同。以 P3 层为例,yolov5s 是 128 通道,yolov5l 是 256 通道,不能直接做 MSE。常见的做法是加一个 1x1 卷积把学生特征投影到老师的通道空间。

为什么不推荐用 3x3?因为 1x1 卷积不改变感受野,只做通道间的线性组合,对齐语义信息更干净,3x3 引入的额外感受野变化会让特征对齐过程学进一些与检测无关的东西。尺寸不一致时用 F.interpolate 插值到老师特征的分辨率,插值模式用 nearest,避免双线性插值带来的平滑效应掩盖特征差异。

import torch import torch.nn as nn import torch.nn.functional as F class FeatureAlignment(nn.Module): def __init__(self, student_channels, teacher_channels): super().__init__() # 1x1 卷积:把学生特征通道数投影到老师通道数 self.align = nn.Conv2d(student_channels, teacher_channels, kernel_size=1, bias=False) def forward(self, student_feat, teacher_feat): student_feat = self.align(student_feat) if student_feat.shape[-2:] != teacher_feat.shape[-2:]: student_feat = F.interpolate( student_feat, size=(teacher_feat.shape[-2], teacher_feat.shape[-1]), mode="nearest", ) return student_feat

特征对齐 loss 本身用 MSE 就够了。MSE 对逐像素差异敏感,在特征图对齐场景下表现稳定,比 L1 收敛更快,也不会像感知 loss 那样需要额外网络。使用的时候,三个尺度分别建一个对齐模块,放在一个 ModuleList 里,跟随学生模型一起训练。

4.2 前景mask与KL蒸馏loss

蒸馏 loss 不能一视同仁地作用在所有空间位置上。yolov5 的检测头特征图里,绝大多数网格是背景,如果对背景位置也做 KL 散度,学生学到的就是「把背景输出得和老师一样」,反而淹没了前景目标的学习信号。

我这里的做法是:训练时用样本自带的 gt 标签构建前景 mask。每个尺度上,先做一次 anchor 匹配,把 gt 匹配到的网格置为 1,其余位置为 0。KL 蒸馏只在 mask 为 1 的前景位置计算分类分支,回归分支则用 SmoothL1 只在前景位置约束。这样既保留 soft label 的类间知识,又不会让背景主导梯度。

import torch import torch.nn as nn import torch.nn.functional as F class DecoupledKD(nn.Module): def __init__(self, temperature=10.0, gate=0.5): super().__init__() self.T = temperature self.gate = gate def forward(self, student_pred, teacher_pred, mask): # student_pred / teacher_pred 形状: (B, 4+1+C, H, W) # 拆出分类分支:去掉前5个通道(x,y,w,h,obj) s_cls = student_pred[:, 5:, :, :] t_cls = teacher_pred[:, 5:, :, :] # 前景区域 KL 蒸馏 t_log_softmax = F.log_softmax(t_cls / self.T, dim=1) s_log_softmax = F.log_softmax(s_cls / self.T, dim=1) kl = F.kl_div(s_log_softmax, t_log_softmax, reduction="none") kl = kl.sum(dim=1, keepdim=True) * mask # 只在前景位置 kl_loss = kl.sum() / (mask.sum() + 1e-6) # 回归分支:只在前景位置做 SmoothL1 s_reg = student_pred[:, :4, :, :] t_reg = teacher_pred[:, :4, :, :] reg = F.smooth_l1_loss(s_reg, t_reg, reduction="none") reg = reg.mean(dim=1, keepdim=True) * mask reg_loss = reg.sum() / (mask.sum() + 1e-6) # 乘 T^2 恢复梯度幅度 return kl_loss * (self.T ** 2) + reg_loss * 2.0

为什么 KL loss 要乘 T 的平方?因为 softmax 之前 logits 除以 T,会让梯度整体缩小约 T 倍,补偿回去才能让蒸馏 loss 和基础 loss 处于同一量级,否则你调 kd-alpha 时的直觉就无效了。回归分支乘 2.0 是经验值,让坐标约束略高于辅助约束,避免学生在框回归上摆烂。

4.3 把蒸馏loss挂进训练循环:改造train和build_targets

蒸馏 loss 怎么挂进 yolov5 的训练流程,是这套源码的关键。yolov5 的模型结构可以按层遍历,detect 层在 model.model 的最后一个位置。前向时,我先跑一次 teacher,再跑一次 student,把两者的 detect 层输出取出来,和训练样本的 target 一起喂给 DecoupledKD。

挂接代码的核心思路如下:

def forward_with_distill(model, teacher, images, targets, kd_loss_fn): # 老师前向:只前向、不回传梯度 with torch.no_grad(): teacher_pred, teacher_feats = teacher(images) # 学生前向 student_pred, student_feats = model(images) # 原始 yolo loss(box + cls + dfl) loss, loss_items = compute_yolo_loss(student_pred, targets) # 构建前景 mask:把每个 gt 映射到对应尺度的网格位置 masks = build_distill_masks(targets, student_pred) # 蒸馏 loss kd_loss = kd_loss_fn(student_pred, teacher_pred, masks) loss = loss + args.kd_alpha * kd_loss return loss, loss_items, kd_loss

这里的 build_distill_masks 是复用 detect 层的 anchor 匹配逻辑,把匹配到的网格位置标记为前景。如果你不想依赖 anchor 匹配,也可以用 IoU 阈值:预测框和 gt 的 IoU 超过 gate 就认定为前景。前者更贴近 yolov5 原生训练逻辑,后者实现简单但对小目标不友好。我一般建议用前一种方案,虽然要多写一点代码,但蒸馏的位置和基础 loss 关注的位置完全一致。

4.4 参数设置参考:T、alpha、feat-lambda 和 gate

参数配置直接决定蒸馏成败,下面这张表可以当作起始配置,再按自己的数据集微调:

参数默认值设置原则
kd-temperature10分类分支软标签温度,温度越高分布越平滑,背景噪声也会增大
kd-alpha0.5蒸馏 loss 与基础 loss 的融合权重,0.5 表示各占一半
feat-lambda0.1特征对齐 loss 权重,通常只有蒸馏 loss 的 1/5~1/10
gate0.5前景判定 IoU 阈值,数据集中小目标多时降到 0.3
epoch 退火最后 20 epochs蒸馏 loss 逐步降为 0,避免学生过度模仿老师的错误输出

提示:feat-lambda 是唯一一个最容易被调崩的参数。特征对齐 loss 权重一旦大于 0.3,学生模型的检测头会被迫去贴合老师的中间语义,任务损失反而下降变慢,最终 mAP 大概率不升反降。

5. 避坑与排查:yolov5蒸馏实战的5个高频翻车点

5.1 翻车点一:蒸馏之后 mAP 不升反降

现象:训练曲线很漂亮,loss 一路下降,但验证集 mAP 比不蒸馏的基线还低 2~3 个点。

原因:最常见的是 feat-lambda 设太大。特征对齐 loss 把学生模型的中间特征强行拉向老师,但你用的学生模型结构小,通道数、感受野都不同,强行对齐的结果是学生把大量容量花在「模仿老师特征」上,检测头自身的分类和回归学习被挤占。

解决:先把 feat-lambda 降到 0.01,或直接关掉特征对齐,只用 KL 蒸馏跑一版。确认 KL 版本有收益后,再逐步加回特征对齐,从 0.01 开始每次翻倍试。记住:特征对齐是加速收敛的辅助项,不是精度主力。

5.2 翻车点二:小目标 mAP 暴跌

现象:整体 mAP 小幅提升,但小目标类别 AP 下降很明显,尤其是车辆、行人这类远处的小物体。

原因:特征对齐时只对齐了 P5 深层特征或直接对齐了所有尺度。P5 感受野大、分辨率低,包含的多是小目标的位置细节。学生网络结构小,在小目标特征上的表达能力本身就弱,再被老师的高质量特征一拉,分布变形,小目标检测反而退化。

解决:特征对齐只做 P3 和 P4,P5 层不对齐,或者给三个尺度分别设置 feat-lambda,P5 的权重直接置零。小目标多的数据集场景,建议 gate 阈值从 0.5 降到 0.3,让更多小目标位置进入蒸馏计算范围。

5.3 翻车点三:温度一调大就 NaN 或梯度爆炸

现象:把 kd-temperature 调到 20 以上,训练到几百步后 loss 突然变成 NaN。

原因:KL 蒸馏乘以 T 的平方会放大梯度,混合精度训练下更容易触发数值上溢出。此外,老师模型的 logits 取值范围在不同类别上差异较大,某些极端置信度的类别 logits 可能到几十,除 T 后 softmax 分母溢出。

解决:把老师模型的分类 logits 做 clamp,限制在 [-10, 10] 区间,再进 softmax。同时 KL 的 T 平方补偿改成 min(T, 10) 的平方,也就是温度超过 10 后不再提高梯度补偿倍数。我之前就吃过这个亏,温度从 8 调到 20,然后整次训练白跑。

5.4 翻车点四:训练速度慢一半、显存翻倍

现象:加上蒸馏后单 step 时间几乎翻倍,显存也翻倍,batch size 被迫减半,最终精度还不如不减半的原始训练。

原因:teacher 前向和 student 前向是两个完整的前向过程,显存自然翻倍。而学生模型前向本身也需要梯度回传,计算结果不能提前释放。batch size 减半后,基础的 batch normalization 统计量变差,学生模型训练不稳定。

解决:老师模型在训练过程中用 fp16 精度做前向,并且不保存中间梯度。更极致的方法是预生成老师特征缓存:先把所有训练图片的老师输出特征存成离线缓存文件,训练时直接读取,省去每次前向老师的时间。这个方案对显存的节约最明显,但需要额外的磁盘空间,我一般只在大数据集上用。

5.5 翻车点五:验证集 mAP 提升是虚的,上线效果差

现象:验证集 mAP 涨了 2 个点,但在真实场景的测试集上反而更差,尤其是有遮挡、光照变化的场景。

原因:老师模型在训练集上过拟合了,软标签里带有偏置,学生把老师的过拟合也学走了。另一种可能是验证集和训练集同分布,蒸馏提升的表面字节来自 teacher 对验证集的记忆。

解决:一是看老师模型的验证集和训练集 loss 差距,差距过大就先纠正老师本身;二是给蒸馏 loss 加退火,训练最后 20 个 epoch 把 kd_alpha 逐步降到 0,让最后阶段回到纯任务损失,把学生从老师的错误记忆里「拉回来」。这也是我后续永久的保留了调度策略,一定要用 epoch 退火。

6. 再往前一步:解耦蒸馏的进阶用法、注意力迁移与效果验证

6.1 进阶一:把蒸馏推进到注意力层面

特征对齐蒸馏的是特征图本身,而注意力迁移蒸馏的是「特征图里哪些位置更重要」这一信息。做法很简单:把老师和学生的特征图分别做一个空间注意力 map,再让两个注意力 map 接近。

def attention_map(feature): # 空间注意力:通道维求平方和再归一化 attn = feature.pow(2).mean(dim=1, keepdim=True) return F.normalize(attn, p=2, dim=2) attention_loss = F.mse_loss( attention_map(student_feat), attention_map(teacher_feat).detach(), )

这个 loss 对特征图内部的值不敏感,只关注哪些位置应该被强调,适合用在学生容量明显不足的场景。它的优点是约束更宽松,学生自由度高;缺点是信息量少,只能作为辅助 loss 使用。我建议把注意力 loss 作为特征对齐模块的补充,直接加在同一个特征对上,权重取 0.05 左右。

6.2 进阶二:验证一套蒸馏效果,不能只看单次 mAP

蒸馏训练有随机性,单次结果不能说明问题。我现在做蒸馏验证的标准动作是三跑取中位数:同一个配置跑三次,每个 epoch 记一次验证 mAP,画三条曲线。如果三条曲线走势接近,结论才可信;如果三条差距超过 0.5 个点,大概率是学习率、batch size 或随机种子的问题,不是蒸馏本身有效。

除了 mAP,还要单独统计「老师错学生也不对」和「老师错学生对」这两类样本的比例。方法很简单,加载老师和蒸馏后学生两个权重,在验证集上跑批量推理,对每个预测结果按 IoU 匹配 gt,统计这两类占比。如果「老师错学生对」的比例能到 10% 以上,说明蒸馏没有把学生变成老师的复制品,而是学到了真正的判别能力。这是很多人评估蒸馏时会忽略的一步。

6.3 调参的主线逻辑与我的收尾习惯

如果你要做多组对比实验,我建议的主线顺序是:基线(纯训练)→ 加 KL 蒸馏 → 加特征对齐 → 加注意力迁移 → 加退火。每一步都在前一步基础上只改一个变量,不要同时调温度和 alpha。这个顺序能让你明确每部分贡献了多少精度,而不是最后 mAP 涨了却说不清是哪里起的作用。

我自己在实战项目里吃过亏:一上来就把所有模块全堆上,loss 花了好几天才降下去,最后还稳不住。后来养成了习惯,任何蒸馏实验先跑最小版本,再把模块一个一个加进去,每个版本保留权重和曲线。这套方法论比任何一个具体参数都更重要,也是我能持续在这个方向复现效果的原因。

希望这次的源码结构和避坑清单能帮你少走一段弯路,也希望你早日在自己的数据集上跑出第一位数的提升。

本文还有配套的精品资源,点击获取

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询