蒸馏模型（Knowledge Distillation）原理和操作详解

最新推荐文章于 2025-05-05 14:56:10 发布

风小子哦

最新推荐文章于 2025-05-05 14:56:10 发布

阅读量1k

点赞数 19

分类专栏：深度学习知识点文章标签：蒸馏教师学生模型 python 人工智能

本文链接：https://blog.youkuaiyun.com/qq_74693232/article/details/146537954

版权

深度学习知识点专栏收录该内容

14 篇文章

订阅专栏

知识蒸馏（Knowledge Distillation，KD）是一种模型压缩技术，其核心思想是让一个小模型（学生模型，Student Model）从一个大模型（教师模型，Teacher Model）中学习知识，使其在轻量化的同时仍保持较高的性能。

1.知识蒸馏的基本概念

在深度学习中，通常有一个大规模、高性能的教师模型（Teacher Model），但由于其计算量较大，难以部署在资源受限的环境（如移动设备或嵌入式系统）上。蒸馏模型的目的是让一个较小的学生模型（Student Model）学习教师模型的知识，从而在减少计算资源的情况下，依然保持较高的准确率。

2.知识蒸馏的核心思想

在普通深度学习训练中，模型的目标是通过交叉熵损失（Cross-Entropy Loss）来学习真实标签（Hard Labels）。
但是，知识蒸馏引入了一种新的学习目标：软标签（Soft Labels），它由教师模型的输出提供。

(1) 传统训练 vs. 知识蒸馏

传统训练： 直接用真实标签训练模型（如 one-hot 形式的分类标签）。
知识蒸馏训练： 额外引入教师模型的输出，使学生模型不仅学习真实标签，还学习教师模型的“知识”（即概率分布）。

例如，在一个分类任务中，教师模型的输出可能是：

        Teacher Model Output:
        Class 1: 0.9
        Class 2: 0.05
        Class 3: 0.03
        Class 4: 0.02

传统训练中，one-hot 标签可能是 [1, 0, 0, 0]，而教师模型的输出则提供了类别之间的关系信息，即哪些类别比较相似（比如 Class 1 和 Class 2 可能较为接近）。

（2）软标签（Soft Labels）

为了让学生模型更好地学习教师模型的知识，引入了温度参数（Temperature, T）来平滑教师模型的输出概率：

其中：是教师模型在类别 i上的logits（未经过Softmax）

T 是温度参数，T 越大，Softmax 输出的概率分布越平滑（也就是说，不同类别的概率差距被缩小）。

例如：

T = 1（默认Softmax）

Class 1: 0.9
                Class 2: 0.05
                Class 3: 0.03
                Class 4: 0.02
        T = 3（更平滑的分布）

                Class 1: 0.6
                Class 2: 0.2
                Class 3: 0.12
                Class 4: 0.08

学生模型在训练时，会同时最小化：

传统的交叉熵损失（对真实标签进行监督）
KL 散度（Kullback-Leibler Divergence，衡量学生和教师输出概率分布的差异）

3.知识蒸馏的损失函数

学生模型的训练目标是最小化以下损失函数：

其中：

Lce是交叉熵损失（针对真实标签）
Lkd是 KL 散度损失（针对教师模型的 Softmax 输出）
α是权重参数，平衡两种损失
T是温度参数

KL 散度公式：

其中和分别是教师模型和学生模型的 Softmax 输出。

在 PyTorch 中，蒸馏损失可以这样实现：

import torch.nn.functional as F

def distillation_loss(student_logits, teacher_logits, labels, alpha=0.5, T=3):
    """
    student_logits: 学生模型的输出 logits
    teacher_logits: 教师模型的输出 logits
    labels: 真实标签
    alpha: 权重参数
    T: 温度参数
    """
    # 计算软标签的 KL 散度损失
    soft_targets = F.softmax(teacher_logits / T, dim=1)
    student_probs = F.log_softmax(student_logits / T, dim=1)
    loss_kd = F.kl_div(student_probs, soft_targets, reduction="batchmean") * (T ** 2)
    
    # 计算交叉熵损失（针对真实标签）
    loss_ce = F.cross_entropy(student_logits, labels)

    # 组合损失
    return alpha * loss_kd + (1 - alpha) * loss_ce