【2】基于 Docker + YOLOv8 环境实现模型蒸馏实战（GTX1660S + Ubuntu22.04）

本文在已有的 YOLOv8 Docker 训练环境上，实现一个完整的“模型蒸馏（Knowledge Distillation）”训练流程。

适合人群：

• 想学习 AI 模型压缩
• 想让小模型速度更快
• 想部署到边缘设备
• 已经能正常训练 YOLOv8
• 希望进一步学习蒸馏、量化、剪枝等优化技术

本文使用：

• Ubuntu 22.04
• Docker
• GTX1660S（6GB）
• PyTorch
• YOLOv8
• coco128 数据集
• Ultralytics YOLOv8 框架

一、什么是模型蒸馏？

模型蒸馏（Knowledge Distillation）本质上是：

用一个“大模型”指导“小模型”学习。

通常：

例如：

我们可以：

• 先训练一个 YOLOv8m
• 再让 YOLOv8n 学习 YOLOv8m 的“知识”

这样：

• 小模型精度会提升
• 推理速度仍然很快
• 更适合边缘部署

二、为什么需要蒸馏？

很多时候：

• GPU 显存不足
• 边缘设备算力差
• Jetson / 工控机 / IPC 部署资源有限

例如 GTX1660S：

• 只有 6GB 显存
• 跑 yolov8x 很吃力
• batch size 容易 OOM

社区里也有很多 YOLOv8 显存不足问题。

因此实际生产中常见方案：

大模型训练 ↓ 模型蒸馏 ↓ 得到高精度小模型 ↓ TensorRT / ONNX 部署

三、实验环境

1. 主机环境

Ubuntu 22.04 Docker NVIDIA Driver CUDA GTX1660S

2. Docker 镜像

nvcr.io/nvidia/pytorch:24.12-py3

该镜像：

• 已内置 CUDA
• 已内置 PyTorch
• 非常适合训练 YOLOv8

PyTorch 本身对 GPU 加速支持非常完善。

3. 启动容器

docker run -it --gpus all \ --shm-size=16g \ --name yolov8-distill \ -v /home/workspace:/workspace \ nvcr.io/nvidia/pytorch:24.12-py3

参数说明：

四、安装 YOLOv8（已安装可跳过）

1. 克隆源码

cd /workspace git clone https://github.com/ultralytics/ultralytics.git cd ultralytics

官方项目：Ultralytics YOLOv8 Github

2. 安装依赖

pip install -e .

检查：

yolo version

五、准备数据集

这里使用官方 coco128。

yolo detect train \ model=yolov8n.pt \ data=coco128.yaml \ epochs=1

首次运行会自动下载数据集。

六、模型蒸馏原理

普通训练：

图片 → Student → Loss → 更新参数

蒸馏训练：

图片 → Teacher ↓ 图片 → Student ↓ 让 Student 模仿 Teacher 输出

蒸馏核心思想：

Student 不仅学习：

• 标签（Ground Truth）

还学习：

• Teacher 的特征
• Teacher 的输出分布
• Teacher 的分类概率

七、先训练 Teacher 模型

蒸馏之前：

必须先有教师模型。

这里使用：

Teacher : YOLOv8m Student : YOLOv8n

1.获取Teacher模型

获取Teacher模型有两种方式，一种是直接使用官网的yolov8m.pt,一种是自己训练yolov8m.pt作为Teacher模型，这里为例演示方便直接使用官网上的yolov8m.pt模型。

八、YOLOv8 蒸馏实现

YOLOv8 官方没有直接提供蒸馏接口。

因此我们自己实现。

1.创建蒸馏训练脚本

创建：

touch distill_train.py

2.完整蒸馏代码

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(), lr=1e-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(f"epoch={epoch} loss={loss.item()}")

九、蒸馏中的关键参数

1. Temperature（温度）

用于软化概率分布。

常见：

temperature = 2~8

越大：

• Teacher 输出越平滑
• Student 更容易学习

2. Alpha

控制：

• Teacher Loss
• Ground Truth Loss

占比。

loss = alpha * kd_loss + (1-alpha) * gt_loss

常见：

alpha = 0.5~0.9

十、蒸馏效果

典型情况：

即：

用小模型获得接近大模型的效果。

十一、为什么蒸馏有效？

因为 Teacher 学到的不仅是标签。

还包括：

• 类别间关系
• 特征表达
• 深层语义信息

例如：

Teacher 可能知道：

狗 ≈ 狼 汽车 ≈ 卡车

这些信息：

普通标签学不到。

十二、进阶蒸馏方式

除了输出蒸馏：

还可以：

1. Feature Distillation

让 Student 学习中间特征图。

Backbone Feature Neck Feature Head Feature

2. Attention Distillation

学习注意力图。

3. Logit Distillation

最经典方式。

学习分类 logits。

十三、蒸馏后的部署优化

蒸馏完成后：

推荐继续：

蒸馏 → ONNX → TensorRT → FP16 → INT8

最终：

• 更小
• 更快
• 更低功耗

十四、YOLOv8 为什么适合蒸馏？

YOLOv8 本身：

• Anchor-Free
• 网络结构清晰
• Backbone/Neck 解耦明显

因此非常适合：

• 剪枝
• 量化
• 蒸馏

YOLOv8 在检测任务上具有较好的速度与精度平衡。

十五、Docker 训练的优势

使用 Docker 的好处：

很多 AI 项目已经大量采用 Docker 化训练环境。

十六、GTX1660S 训练建议

1660S 只有 6GB 显存。

建议：

后续进阶方向

你后面可以继续学习：

1. Feature Map 蒸馏
2. YOLOv8 剪枝
3. TensorRT INT8 校准
4. PTQ / QAT 量化
5. DeepSparse
6. ONNX Runtime
7. NCNN 部署
8. Jetson 边缘部署