怎样在Python中利用TensorFlow 2.x实现知识蒸馏-Python教程-PHP中文网

知识蒸馏的核心是让学生模型学习教师模型经温度缩放后的logits分布，而非简单复制模型或特征；需自定义混合损失函数，同步计算带温度补偿的kl散度与硬标签交叉熵，并在自定义训练循环中同批数据联合推理。

怎样在python中利用tensorflow 2.x实现知识蒸馏

知识蒸馏的核心不是模型复制，而是让小模型学“ logits 分布”

知识蒸馏在 TensorFlow 2.x 中不依赖特殊 API，关键在于修改损失函数：用教师模型的软标签（经温度缩放的 softmax 输出）指导学生模型。硬标签（真实类别）只起辅助作用。直接复用 model.fit() 会失败——因为默认只算交叉熵，没接入教师输出。

常见错误是把教师模型当固定特征提取器，只蒸馏中间层；这其实是特征迁移，不是标准知识蒸馏。真正的蒸馏必须对齐 logits 层输出，并引入温度 T 控制分布平滑度。

教师模型需预先训练好，且推理时保持 training=False
学生模型输入与教师完全一致（同尺寸、同归一化）
温度 T 通常设为 3–20；T=1 等价于普通交叉熵，失去蒸馏意义
软损失权重常设为 0.7～0.9，硬损失补足剩余部分

自定义蒸馏损失函数必须同时处理软/硬目标

TensorFlow 2.x 推荐用函数式写法实现混合损失，避免 Keras 内置损失无法接入教师 logits 的问题。不能只写 tf.keras.losses.categorical_crossentropy，要手动计算带温度的 KL 散度。

注意：KL 散度要求输入是概率分布（和为 1），所以必须先对教师和学生 logits 做 softmax，且学生侧要除以温度再 softmax，教师侧也需同样温度缩放——否则分布不对齐，梯度爆炸很常见。

def distillation_loss(y_true, y_pred, y_teacher, T=4, alpha=0.8):
    # 学生软目标：logits / T → softmax
    soft_pred = tf.nn.softmax(y_pred / T, axis=-1)
    # 教师软目标：同温度缩放
    soft_teacher = tf.nn.softmax(y_teacher / T, axis=-1)
    # KL 散度（batch-wise，非 sample-wise）
    kl_loss = tf.keras.losses.KLDivergence()(
        soft_teacher, soft_pred
    ) * (T ** 2)  # 温度补偿项，不可省
    # 硬标签交叉熵
    ce_loss = tf.keras.losses.sparse_categorical_crossentropy(
        y_true, y_pred, from_logits=True
    )
    return alpha * kl_loss + (1 - alpha) * ce_loss

这里 T ** 2 是标准补偿因子，省略会导致 KL 损失过小，蒸馏失效；from_logits=True 保证硬损失仍作用于原始 logits，而非 softmax 后结果。