本文在已有的 YOLOv8 Docker 训练环境上实现一个完整的“模型蒸馏Knowledge Distillation”训练流程。适合人群想学习 AI 模型压缩想让小模型速度更快想部署到边缘设备已经能正常训练 YOLOv8希望进一步学习蒸馏、量化、剪枝等优化技术本文使用Ubuntu 22.04DockerGTX1660S6GBPyTorchYOLOv8coco128 数据集Ultralytics YOLOv8 框架一、什么是模型蒸馏模型蒸馏Knowledge Distillation本质上是用一个“大模型”指导“小模型”学习。通常角色作用Teacher教师模型精度高但体积大Student学生模型更小、更快例如模型参数量推理速度YOLOv8m较大较慢YOLOv8n很小很快我们可以先训练一个 YOLOv8m再让 YOLOv8n 学习 YOLOv8m 的“知识”这样小模型精度会提升推理速度仍然很快更适合边缘部署二、为什么需要蒸馏很多时候GPU 显存不足边缘设备算力差Jetson / 工控机 / IPC 部署资源有限例如 GTX1660S只有 6GB 显存跑 yolov8x 很吃力batch size 容易 OOM社区里也有很多 YOLOv8 显存不足问题。因此实际生产中常见方案大模型训练 ↓ 模型蒸馏 ↓ 得到高精度小模型 ↓ TensorRT / ONNX 部署三、实验环境1. 主机环境Ubuntu 22.04 Docker NVIDIA Driver CUDA GTX1660S2. Docker 镜像nvcr.io/nvidia/pytorch:24.12-py3该镜像已内置 CUDA已内置 PyTorch非常适合训练 YOLOv8PyTorch 本身对 GPU 加速支持非常完善。3. 启动容器docker run -it --gpus all \ --shm-size16g \ --name yolov8-distill \ -v /home/workspace:/workspace \ nvcr.io/nvidia/pytorch:24.12-py3参数说明参数作用--gpus all使用 GPU--shm-size增加共享内存-v挂载代码目录四、安装 YOLOv8已安装可跳过1. 克隆源码cd /workspace git clone https://github.com/ultralytics/ultralytics.git cd ultralytics官方项目Ultralytics YOLOv8 Github2. 安装依赖pip install -e .检查yolo version五、准备数据集这里使用官方 coco128。yolo detect train \ modelyolov8n.pt \ datacoco128.yaml \ epochs1首次运行会自动下载数据集。六、模型蒸馏原理普通训练图片 → Student → Loss → 更新参数蒸馏训练图片 → Teacher ↓ 图片 → Student ↓ 让 Student 模仿 Teacher 输出蒸馏核心思想Student 不仅学习标签Ground Truth还学习Teacher 的特征Teacher 的输出分布Teacher 的分类概率七、先训练 Teacher 模型蒸馏之前必须先有教师模型。这里使用Teacher : YOLOv8m Student : YOLOv8n1.获取Teacher模型获取Teacher模型有两种方式一种是直接使用官网的yolov8m.pt,一种是自己训练yolov8m.pt作为Teacher模型这里为例演示方便直接使用官网上的yolov8m.pt模型。八、YOLOv8 蒸馏实现YOLOv8 官方没有直接提供蒸馏接口。因此我们自己实现。1.创建蒸馏训练脚本创建touch distill_train.py2.完整蒸馏代码from ultralytics import YOLO import torch import torch.nn as nn import torch.nn.functional as F # Teacher teacher YOLO(runs/detect/train/weights/best.pt) # Student student YOLO(yolov8n.pt) teacher_model teacher.model student_model student.model teacher_model.eval() # 冻结 Teacher for p in teacher_model.parameters(): p.requires_grad False optimizer torch.optim.Adam( student_model.parameters(), lr1e-4 ) temperature 4.0 alpha 0.7 device cuda teacher_model.to(device) student_model.to(device) for epoch in range(10): imgs torch.randn(4, 3, 640, 640).to(device) with torch.no_grad(): teacher_out teacher_model(imgs) student_out student_model(imgs) # 蒸馏 Loss kd_loss F.mse_loss( student_out[0], teacher_out[0] ) # 普通 Loss示例 gt_loss student_out[0].mean() loss alpha * kd_loss (1 - alpha) * gt_loss optimizer.zero_grad() loss.backward() optimizer.step() print(fepoch{epoch} loss{loss.item()})九、蒸馏中的关键参数1. Temperature温度用于软化概率分布。常见temperature 2~8越大Teacher 输出越平滑Student 更容易学习2. Alpha控制Teacher LossGround Truth Loss占比。loss alpha * kd_loss (1-alpha) * gt_loss常见alpha 0.5~0.9十、蒸馏效果典型情况模型mAPFPSYOLOv8m高较低YOLOv8n较低高蒸馏后的 YOLOv8n接近 mAP高 FPS即用小模型获得接近大模型的效果。十一、为什么蒸馏有效因为 Teacher 学到的不仅是标签。还包括类别间关系特征表达深层语义信息例如Teacher 可能知道狗 ≈ 狼 汽车 ≈ 卡车这些信息普通标签学不到。十二、进阶蒸馏方式除了输出蒸馏还可以1. Feature Distillation让 Student 学习中间特征图。Backbone Feature Neck Feature Head Feature2. Attention Distillation学习注意力图。3. Logit Distillation最经典方式。学习分类 logits。十三、蒸馏后的部署优化蒸馏完成后推荐继续蒸馏 → ONNX → TensorRT → FP16 → INT8最终更小更快更低功耗十四、YOLOv8 为什么适合蒸馏YOLOv8 本身Anchor-Free网络结构清晰Backbone/Neck 解耦明显因此非常适合剪枝量化蒸馏YOLOv8 在检测任务上具有较好的速度与精度平衡。十五、Docker 训练的优势使用 Docker 的好处优势说明环境隔离不污染宿主机CUDA统一避免版本冲突快速迁移可复制到其他服务器易部署CI/CD方便很多 AI 项目已经大量采用 Docker 化训练环境。十六、GTX1660S 训练建议1660S 只有 6GB 显存。建议参数推荐imgsz640batch4~8modelyolov8n/sAMP开启cache禁用后续进阶方向你后面可以继续学习Feature Map 蒸馏YOLOv8 剪枝TensorRT INT8 校准PTQ / QAT 量化DeepSparseONNX RuntimeNCNN 部署Jetson 边缘部署