简介:这份资源面向希望将知识蒸馏落地到目标检测任务的算法工程师与深度学习学习者,聚焦 yolov5 模型轻量化场景。知识蒸馏通过教师网络指导学生网络训练,把大模型学到的知识迁移到小模型上,从而在压缩参数量的同时尽量保持检测精度,是模型部署与加速的常用手段。资源包共 8 个文件,以 4 个 zip 压缩包、3 个 pt 权重文件和 1 个 py 脚本为主,整体约 593MB,其中压缩包分别承载蒸馏代码、数据集与说明文档,pt 文件为可直接加载的预训练权重,py 脚本负责数据准备与处理。内容覆盖完整知识蒸馏原理教程、环境配置说明、数据处理代码以及 yolov5 蒸馏实现,按教程操作即可理解蒸馏流程并掌握代码细节。目前已有 2022 人学习下载,适合想系统掌握检测模型蒸馏与轻量化改造的读者参考实践。
1. 基于 YOLOv5 的知识蒸馏:为什么教师模型不是越大越好
你训好了一个 YOLOv5l,mAP 比 YOLOv5s 高了 4 个点,但部署到树莓派 5 或者边缘盒子上,推理一帧要 200ms 以上,产线根本等不起。这时候大多数人的第一反应是「换个更小的模型重新训」,但小模型从零训往往收敛到次优解,精度掉得心疼。知识蒸馏要解决的就是这个矛盾:让大模型当教师,把它的「暗知识」——那些不是硬标签、而是软概率分布里藏着的类间相似性——迁移到小模型身上,让小模型在参数量不变的前提下逼近大模型的精度。
这个方向适合两类人:一类是手里已经有 YOLOv5 训练管线、想在不换硬件的前提下榨出几个点 mAP 的工程师;另一类是想搞懂蒸馏到底怎么落到检测任务上、而不是只停留在分类 MNIST 例子上的人。源码层面,YOLOv5 的蒸馏不是官方内置功能,需要自己在训练循环里插损失项,这也是为什么「基于 yolov5 的知识蒸馏实战源码」这个标题下,真正能跑通的东西不多——大部分文章只讲概念,不给可复现的 loss 拼接和特征对齐代码。接下来我按自己踩过的路子,把教师学生怎么选、损失怎么拼、特征层怎么对齐、坑在哪,一层层拆开。
2. 蒸馏前的选型:教师学生配对与 YOLOv5 结构对齐点
2.1 教师和学生的参数量差距控制在多少才有效
知识蒸馏不是教师越大越好。我试过用 YOLOv5x 教 YOLOv5n,mAP 反而比 YOLOv5s 教 YOLOv5n 低了 0.8 个点。原因不玄学:教师和学生容量差距过大时,学生的拟合能力根本吃不下教师的软分布,蒸馏损失降不下去,反而干扰了原本的检测损失。常见做法是教师比学生大一到两个量级,比如 YOLOv5m 教 YOLOv5s、YOLOv5l 教 YOLOv5m。如果你非要用 x 教 n,至少把温度 T 调高到 4 以上,让软标签更平滑,降低学生拟合难度。
另一个选型点是教师是否已经在自己的数据集上微调过。直接拿 COCO 预训练权重当教师,软标签里全是 COCO 的类间关系,跟你自己的安全帽、缺陷检测数据集分布不匹配,蒸馏收益会打对折。我一般会先把教师在自己的训练集上 fine-tune 到收敛,再冻结它当教师。这一步多花两小时,但学生最终 mAP 能多拿 1.5 个点以上。
2.2 YOLOv5 里哪几层适合做特征蒸馏
YOLOv5 的 backbone 是 CSPDarknet,neck 是 PANet,head 是三个不同尺度的检测头。特征蒸馏不是随便挑一层对齐就行。分类任务蒸馏通常只蒸 logits,但检测任务里,head 输出的 logits 包含分类分支和回归分支,直接蒸会互相干扰。我一般选三个位置:
| 蒸馏位置 | 层名示例 | 蒸馏内容 | 权重建议 |
|---|---|---|---|
| Neck 输出 P3 | model.model[15] | 特征图 MSE | 0.5 |
| Neck 输出 P4 | model.model[18] | 特征图 MSE | 0.5 |
| Neck 输出 P5 | model.model[21] | 特征图 MSE | 0.5 |
| Head 分类分支 | model.model[24].cv3 | 软标签 KL | 1.0 |
P3/P4/P5 是 PANet 融合后的多尺度特征,语义信息比 backbone 中间层强,又比 head 输出干净。学生和教师的这三个特征图空间尺寸一致(因为输入分辨率相同),通道数不同,所以对齐时需要 1x1 卷积做通道映射。head 分类分支的软标签蒸馏用 KL 散度,回归分支不蒸,因为框坐标的软分布没有明确物理意义,蒸了容易让学生的框回归变糊。
2.3 蒸馏温度 T 和损失权重的初始值怎么定
温度 T 控制软标签的平滑程度。T=1 就是原始 softmax,T 越大,负类概率被放大越多,暗知识越明显。检测任务里我一般从 T=3 起步,如果学生收敛慢就调到 4,如果学生 mAP 抖动大就降到 2。损失权重方面,总损失是:
L_total = L_det + alpha * L_feat + beta * L_logitL_det 是 YOLOv5 原本的框回归+分类+置信度损失。alpha 控制特征蒸馏权重,beta 控制 logit 蒸馏权重。我的经验初始值:alpha=0.5,beta=1.0。如果学生 mAP 比不蒸馏还低,先把 beta 降到 0.5 试,再不行就检查教师是否过拟合。alpha 超过 1.0 容易让学生特征图过度模仿教师,反而丢失自己从数据里学到的细节。
3. 把蒸馏损失接进 YOLOv5 训练循环:源码级改动
3.1 教师模型加载与冻结的正确姿势
YOLOv5 的 train.py 里,模型构建在Model类里完成。要加蒸馏,最干净的做法是新建一个DistillModel包装类,而不是直接改Model的 forward。下面是我用的教师加载和冻结代码:
import torch import torch.nn as nn class DistillModel(nn.Module): def __init__(self, student, teacher, cfg): super().__init__() self.student = student self.teacher = teacher self.cfg = cfg # 教师冻结,不参与梯度更新 for p in self.teacher.parameters(): p.requires_grad = False self.teacher.eval() # 通道映射层:学生通道 -> 教师通道 # 假设学生 P3/P4/P5 通道为 [128, 256, 512] # 教师 P3/P4/P5 通道为 [256, 512, 1024] self.align_p3 = nn.Conv2d(128, 256, 1) self.align_p4 = nn.Conv2d(256, 512, 1) self.align_p5 = nn.Conv2d(512, 1024, 1) def forward(self, x): # 学生前向,拿中间特征 student_feats = self.student.forward_features(x) with torch.no_grad(): teacher_feats = self.teacher.forward_features(x) return student_feats, teacher_feats逻辑说明:教师用torch.no_grad()包住前向,避免显存浪费在教师梯度上。forward_features需要你在 YOLOv5 的Model类里加一个方法,返回 P3/P4/P5 三层特征和 head 分类输出。通道映射层用 1x1 卷积,参数量极小,不会拖慢训练。注意教师要.eval(),因为 YOLOv5 里 BN 层在 train 模式下会更新 running stats,教师必须保持推理模式。
参数说明:cfg里放温度 T、alpha、beta。通道数 128/256/512 是 YOLOv5s 的典型值,如果你用 YOLOv5m 或 l,需要打印model.model[15].cv2.conv.out_channels确认实际通道数再改。
3.2 特征蒸馏损失和 logit 蒸馏损失的拼接
损失计算放在训练循环里,每个 batch 前向之后调用:
import torch.nn.functional as F def distill_loss(student_feats, teacher_feats, align_layers, T, alpha, beta): # 特征蒸馏:MSE + 通道对齐 s_p3, s_p4, s_p5, s_cls = student_feats t_p3, t_p4, t_p5, t_cls = teacher_feats loss_feat = 0.0 for s_feat, t_feat, align in zip( [s_p3, s_p4, s_p5], [t_p3, t_p4, t_p5], [align_layers[0], align_layers[1], align_layers[2]] ): s_aligned = align(s_feat) loss_feat += F.mse_loss(s_aligned, t_feat.detach()) loss_feat = loss_feat / 3.0 # logit 蒸馏:KL 散度,带温度 # s_cls/t_cls 形状 [B, num_anchors, num_classes] s_logit = F.log_softmax(s_cls / T, dim=-1) t_prob = F.softmax(t_cls / T, dim=-1) loss_logit = F.kl_div(s_logit, t_prob, reduction='batchmean') * (T * T) return alpha * loss_feat + beta * loss_logit逻辑说明:特征蒸馏用 MSE,教师特征.detach()切断梯度,只让学生往教师方向靠。logit 蒸馏用 KL 散度,T*T是标准蒸馏里的梯度缩放因子,因为 softmax 除以 T 后梯度会缩小 T² 倍,乘回来保持量级。reduction='batchmean'比'mean'更稳定,因为 anchor 数量多,mean 会把损失压得太小。
参数说明:T 从 3 开始调。alpha 和 beta 初始 0.5 和 1.0。如果显存不够,把特征蒸馏只保留 P4/P5,去掉 P3,因为 P3 特征图最大,MSE 显存占用最高。
3.3 训练循环里怎么把蒸馏损失和检测损失相加
YOLOv5 的ComputeLoss类返回检测损失。你需要在train.py里拿到检测损失后,加上蒸馏损失再 backward:
# 在 train.py 的训练循环内 pred = distill_model(x) # 检测损失,用学生输出算 det_loss, det_items = compute_loss(pred_student, targets) # 蒸馏损失 d_loss = distill_loss( student_feats, teacher_feats, [distill_model.align_p3, distill_model.align_p4, distill_model.align_p5], T=3, alpha=0.5, beta=1.0 ) total_loss = det_loss + d_loss total_loss.backward()逻辑说明:检测损失只从学生输出算,教师不参与。蒸馏损失加在总损失里,梯度只回传到学生和通道映射层。注意compute_loss的输入必须是学生 head 输出,不能混入教师输出。如果你用 YOLOv5 的--hyp超参文件,可以把 alpha/beta/T 写进 hyp.yaml,在 train.py 里读出来,方便做超参搜索。
参数说明:T=3是温度,alpha=0.5是特征蒸馏权重,beta=1.0是 logit 蒸馏权重。这三个值建议先用默认跑 10 个 epoch,看 mAP 曲线再调。
4. 蒸馏训练避坑:从 NaN 损失到 mAP 不升反降
4.1 损失出现 NaN:温度 T 和 log_softmax 的数值陷阱
现象:训练到第 3 个 epoch,total_loss 突然变成 NaN,之后所有梯度都是 NaN。
原因:KL 散度里t_prob经过 softmax 后可能出现 0 值,log_softmax输出负无穷,乘起来 NaN。教师 logit 如果本身有极端值(比如某类 logit 到 50),除以 T=3 后仍然很大,softmax 后其他类概率被压到 0。
解决:在t_prob上加一个极小值1e-8,并且对教师 logit 做 clamp:
t_cls = torch.clamp(t_cls, -20, 20) t_prob = F.softmax(t_cls / T, dim=-1) + 1e-8 t_prob = t_prob / t_prob.sum(dim=-1, keepdim=True)另外检查教师是否在 train 模式,BN 统计量漂移也会导致 logit 爆炸。
4.2 学生 mAP 比不蒸馏还低:教师过拟合或权重失衡
现象:蒸馏训练 50 epoch,学生 mAP 比 baseline 低 1.2 个点。
原因:教师在自己的训练集上过拟合,软标签里全是训练集噪声,学生学了反而泛化差。或者 beta 太大,学生过度模仿教师分类分布,丢失了从硬标签学到的判别边界。
解决:先验证教师在自己验证集上的 mAP,如果教师验证 mAP 比训练 mAP 低超过 5 个点,说明过拟合,需要先给教师加 dropout 或早停。然后降 beta 到 0.3,升 alpha 到 0.8,让特征蒸馏占主导。我遇到过最坑的一次是教师用了 mosaic 增强,学生没用,软标签里的增强噪声直接带偏学生。
4.3 显存爆掉:特征蒸馏的 MSE 比检测损失更吃显存
现象:batch size 从 16 降到 8 还是 OOM,但 baseline 用 16 没问题。
原因:P3 特征图尺寸是 80x80,通道 128,MSE 计算时align(s_feat)和t_feat都要存中间激活,显存占用是检测损失的 2 倍以上。
解决:特征蒸馏只保留 P4/P5,去掉 P3。或者用torch.utils.checkpoint对 align 层做梯度检查点。再不行就把特征蒸馏的 batch 拆成两个半 batch 分别算 MSE 再平均。我一般直接砍 P3,mAP 损失不到 0.3 个点,显存省 30%。
4.4 教师学生输入分辨率不一致导致特征图对不上
现象:F.mse_loss报错 shape mismatch,学生 P3 是 80x80,教师 P3 是 40x40。
原因:YOLOv5 的--img-size对教师和学生用了不同值,或者教师加载时没同步imgsz参数。
解决:在DistillModel初始化时强制教师和学生的imgsz一致,并且检查model.stride是否相同。如果教师是 YOLOv5l,学生是 YOLOv5s,stride 都是 32/16/8,只要输入分辨率一样,特征图尺寸就一样。别在蒸馏时开--rect矩形训练,不同 batch 的尺寸变化会让特征图尺寸动态变,MSE 直接崩。
4.5 蒸馏后模型导出 ONNX 失败:自定义层没注册
现象:训练完导出 ONNX,报Unsupported operator: align_p3。
原因:通道映射层是自定义 Conv2d,导出时 PyTorch 能识别 Conv2d,但如果你用了自定义的AlignLayer类,ONNX 不认。
解决:通道映射直接用nn.Conv2d,不要包自定义类。导出前把DistillModel拆开,只导出student,教师和 align 层不导出。YOLOv5 的export.py里指定model=distill_model.student即可。
5. 蒸馏效果验证与进阶:用 mAP 曲线和特征可视化判断值不值得做
蒸馏训练跑完,怎么判断它真的有效,而不是随机波动?我一般看三条线:学生 mAP@0.5、学生 mAP@0.5:0.95、以及学生和教师特征图的 CKA 相似度。mAP 曲线如果在前 20 epoch 就超过 baseline,说明蒸馏信号有效;如果 50 epoch 才勉强追平,说明权重没调好。CKA 相似度用torchmetrics里的CKA算 P4 特征,蒸馏后学生和教师的 CKA 应该从 0.4 左右升到 0.7 以上,如果没升,说明特征蒸馏没起作用。
进阶用法有两个方向。一是自蒸馏:用同一个模型的不同深度特征互相蒸,比如用 YOLOv5s 的 P5 蒸 P3,不需要额外教师,显存省一半。二是在线蒸馏:教师和学生同时更新,教师用学生的 EMA 权重,适合教师也没训好的场景。我试过在线蒸馏在安全帽数据集上比离线蒸馏多拿 0.6 个点,但训练不稳定,需要把教师更新动量设到 0.999 以上。
验证蒸馏是否值得做,最直接的办法是固定学生结构,跑三组:baseline 不蒸馏、只蒸 logit、特征+logit 都蒸。如果特征+logit 比 baseline 高不到 1 个点,而训练时间多了 40%,那就不值得,不如直接换 YOLOv5m。我自己的习惯是:边缘设备上 YOLOv5s 蒸馏后能到 YOLOv5m 的 95% 精度,但推理速度还是 s 的速度,这时候才值得投入。如果蒸馏后只比 s 高 0.5 个点,我宁愿用 s 加更多数据增强。希望帮到你。
本文还有配套的精品资源,点击获取