TF示例演示教师模型训练、学生模型蒸馏与结果导出

一、说明

本文件代码以下记录了一个完整、可运行的 TensorFlow 示例，演示“知识蒸馏”是怎么从教师模型训练、到学生模型蒸馏、再到模型导出与结果可视化的。

• 本文件代码主线作用

  • 加载并预处理 FashionMNIST 数据集，让图片和标签变成适合训练的形式： 数据加载

  • 构建一个较大的 教师模型 ，先把它单独训练好： 教师模型定义 、 教师模型训练

  • 构建一个更小的 学生模型 ，让它不直接只学真实标签，而是去学习教师模型输出的“软目标”： 学生模型定义 、 蒸馏损失

  • 用自定义训练循环完成学生模型的蒸馏训练，包括梯度计算、参数更新、loss/acc 统计： 训练与验证流程

  • 最后把教师模型和学生模型转成 TFLite ，并画出训练曲线，方便比较效果和后续部署： TFLite 导出 、 损失曲线绘制

• 本文件代码解决什么问题

  • 大模型通常效果更好，但参数更多、部署更重

  • 小模型更轻量，但直接训练时效果可能差一些

  • 知识蒸馏就是让小模型向大模型“学经验”，争取在更小体积下保留尽量好的效果

二、简洁注释版文件代码

# -*- coding: utf-8 -*-
"""
此文件代码已经上传保存到google colab上，可以直接运行。
colab地址为：https://colab.research.google.com/drive/1sFxDyFvHRZr5YvV8S2PFw-583RaucKS3
运行的主要环境包括：Python 3.12.13、tensorflow 2.20.0
"""

# 导入必要的库
import os
import tensorflow as tf

from tensorflow.keras import models
from tensorflow.keras import layers

# 设置随机种子以保证结果可复现
tf.random.set_seed(666)

# 加载FashionMNIST数据集，并对像素值进行缩放
(X_train, y_train), (X_test, y_test) = tf.keras.datasets.fashion_mnist.load_data()

# 将训练集像素值缩放到0-1之间
X_train = X_train/255.
# 将测试集像素值缩放到0-1之间
X_test = X_test/255.

# 打印训练集和测试集的形状
X_train.shape, X_test.shape, y_train.shape, y_test.shape

# 将像素值转换为float32类型，并重塑输入数据的形状
X_train = X_train.astype("float32").reshape(-1, 28, 28, 1)
X_test = X_test.astype("float32").reshape(-1, 28, 28, 1)

# 定义构建基本浅层卷积神经网络的工具函数
def get_teacher_model():
    # 初始化顺序模型
    model = models.Sequential()
	# 添加卷积层，16个卷积核，大小5x5，ReLU激活
	model.add(layers.Conv2D(16, (5, 5), activation="relu",
        input_shape=(28, 28, 1)))
    # 添加最大池化层，池化窗口2x2
    model.add(layers.MaxPooling2D(pool_size=(2, 2)))
    # 再添加一个卷积层，32个卷积核，大小5x5，ReLU激活
    model.add(layers.Conv2D(32, (5, 5), activation="relu"))
    # 再添加一个最大池化层，池化窗口2x2
    model.add(layers.MaxPooling2D(pool_size=(2, 2)))
    # 添加Dropout层，丢弃率为0.2，防止过拟合
    model.add(layers.Dropout(0.2))
    # 将多维输出展平为一维
    model.add(layers.Flatten())
    # 添加全连接层，128个神经元，ReLU激活
    model.add(layers.Dense(128, activation="relu"))
    # 添加输出层，10个神经元（对应10个类别）
    model.add(layers.Dense(10))

	# 返回构建好的模型
    return model

# 定义损失函数和优化器
loss_func = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
# 使用Adam优化器
optimizer = tf.keras.optimizers.Adam()

# 准备TensorFlow数据集
# 创建训练数据管道：从tensor切片创建dataset，打乱缓冲区大小100，批次大小64
train_ds = tf.data.Dataset.from_tensor_slices((X_train, y_train)).shuffle(100).batch(64)
# 创建测试数据管道：从tensor切片创建dataset，批次大小64
test_ds = tf.data.Dataset.from_tensor_slices((X_test, y_test)).batch(64)

# 训练教师模型
# 获取教师模型实例
teacher_model = get_teacher_model()

# 编译模型，指定损失函数、优化器和评估指标
teacher_model.compile(loss=loss_func, optimizer=optimizer, metrics=["accuracy"])
# 训练模型，使用 train_ds 作为训练数据， test_ds 作为验证数据，训练10个轮次
history = teacher_model.fit(train_ds,
                  validation_data=test_ds,
                  epochs=10)

# 评估模型并保存权重
# 打印在测试集上的准确率，保留两位小数
print("Test accuracy: {:.2f}".format(teacher_model.evaluate(test_ds)[1]*100))
# 保存教师模型权重到文件
teacher_model.save_weights("teacher_model.weights.h5")

# 学生模型工具函数
def get_student_model():
    # 初始化顺序模型
    model = models.Sequential()
    # 添加输入层，形状为(28, 28, 1)
    model.add(layers.Input(shape=(28, 28, 1)))
    # 将输入展平
    model.add(layers.Flatten())
    # 添加全连接层，48个神经元，ReLU激活
    model.add(layers.Dense(48, activation="relu"))
    # 添加输出层，10个神经元
    model.add(layers.Dense(10))

    # 返回构建好的模型
    return model

# 来源：https://github.com/google-research/simclr/blob/master/colabs/distillation_self_training.ipynb
def get_kd_loss(student_logits, teacher_logits, temperature=0.5):
    """
    知识蒸馏的核心损失函数。
    作用：计算学生模型和教师模型在“软目标（Soft Targets）”上的差异。
    """

    # 计算教师模型的软标签概率，使用温度参数缩放logits
    teacher_probs = tf.nn.softmax(teacher_logits / temperature)
    # 计算知识蒸馏损失（使用兼容v1版本的loss函数，计算softmax交叉熵）
    # 注意这里student_logits也除了温度，并且损失乘以温度的平方以保持梯度量级
    kd_loss = tf.compat.v1.losses.softmax_cross_entropy(teacher_probs, student_logits / temperature, temperature**2)
    return kd_loss

# 模型和优化器
# 获取学生模型实例
student_model = get_student_model()
# 定义优化器，学习率为0.01
optimizer = tf.keras.optimizers.Adam(learning_rate=0.01)

# 计算每个epoch内batch的平均损失
train_loss = tf.keras.metrics.Mean(name="train_loss")
valid_loss = tf.keras.metrics.Mean(name="test_loss")

# 指定性能指标
train_acc = tf.keras.metrics.SparseCategoricalAccuracy(name="train_acc")
valid_acc = tf.keras.metrics.SparseCategoricalAccuracy(name="valid_acc")


# 训练工具函数
@tf.function
def model_train(images, labels, teacher_model,
                student_model, optimizer, temperature):
    # 获取教师模型对图像的预测logits（不进行梯度追踪）
    teacher_logits = teacher_model(images)

    # 开启梯度记录
    with tf.GradientTape() as tape:
        # 获取学生模型对图像的预测logits
        student_logits = student_model(images)
        # 计算知识蒸馏损失
        loss = get_kd_loss(student_logits, teacher_logits, temperature)

        # 计算相对于学生模型可训练变量的梯度
        gradients = tape.gradient(loss, student_model.trainable_variables)
        # 应用梯度更新学生模型的参数
        optimizer.apply_gradients(zip(gradients, student_model.trainable_variables))

        # 更新训练损失指标
        train_loss(loss)
        # 更新训练准确率指标
        train_acc(labels, tf.nn.softmax(student_logits))


# 验证工具函数
@tf.function
def model_validate(images, labels, teacher_model,
                   student_model, temperature):
    # 获取教师模型预测logits
    teacher_logits = teacher_model(images)

    # 获取学生模型预测logits
    student_logits = student_model(images)
    # 计算知识蒸馏损失
    loss = get_kd_loss(student_logits, teacher_logits, temperature)

    # 更新验证损失指标
    valid_loss(loss)
    # 更新验证准确率指标
    valid_acc(labels, tf.nn.softmax(student_logits))


# 将所有步骤结合在一起进行训练
def train_model(epochs, teacher_model, student_model, optimizer, temperature=0.5):
    # 遍历每个epoch
    for epoch in range(epochs):
        # 遍历训练数据集中的每个batch
        # 这里的 images, labels 长度通常为 64 (因为我们在定义 train_ds 时设置了 .batch(64))
        # 唯一的例外是最后一个 batch，如果数据总数无法整除 64，剩下的余数会在最后一次取出。
        for (images, labels) in train_ds:
            # 执行单步训练
            model_train(images, labels, teacher_model, student_model, optimizer, temperature)

        # 遍历测试数据集中的每个batch
        for (images, labels) in test_ds:
            # 执行单步验证
            model_validate(images, labels, teacher_model, student_model, temperature)

        # 获取当前epoch的平均损失和准确率
        (loss, acc) = train_loss.result(), train_acc.result()
        (val_loss, val_acc) = valid_loss.result(), valid_acc.result()

        # 重置指标状态，为下一个epoch做准备
        train_loss.reset_state(), train_acc.reset_state()
        valid_loss.reset_state(), valid_acc.reset_state()

        # 定义打印模板
        template = "Epoch {}, loss: {:.3f}, acc: {:.3f}, val_loss: {:.3f}, val_acc: {:.3f}"
        # 打印当前epoch的训练结果
        print (template.format(epoch+1,
                            loss,
                            acc,
                            val_loss,
                            val_acc))

    # 返回训练好的教师和学生模型
    return teacher_model, student_model

# 开始训练模型，训练10个epoch
_, student_model = train_model(10, teacher_model, student_model, optimizer)
# 可以通过更长的训练时间和更仔细的超参数调整进一步改进。

# 序列化保存
# 保存学生模型权重
student_model.save_weights("student_model.weights.h5")

# 查看模型文件大小
# 列出当前目录下所有的h5文件及其详细信息，将可以看到学生模型比教师模型在效果比较接近的情况下，体积要小得多（结构也要简单很多）
os.system("ls -lh *.h5")

# 打印教师模型摘要
# 通过摘要查看教师模型的层结构、输出形状和参数数量
teacher_model.summary()

# 打印学生模型摘要
# 通过摘要查看学生模型的层结构、输出形状和参数数量
student_model.summary()

# 使用TFLite可以进一步减小模型大小。
# 来源：https://www.tensorflow.org/lite/performance/post_training_quant

# 定义代表性数据集生成器函数
def representative_data_gen():
    # 从训练数据中取100个样本，每次一个，用于量化校准
    for input_value in tf.data.Dataset.from_tensor_slices(X_train).batch(1).take(100):
        yield [input_value]

# 定义转换为TFLite模型的函数
def convert_to_tflite(model, tflite_file):
    # 从Keras模型创建TFLite转换器
    converter = tf.lite.TFLiteConverter.from_keras_model(model)
    # 设置优化选项为默认优化
    converter.optimizations = [tf.lite.Optimize.DEFAULT]
    # 设置代表性数据集用于量化
    converter.representative_dataset = representative_data_gen
    # 设置目标规范支持的操作集为INT8内置操作
    converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
    # 设置推理输入类型为int8
    converter.inference_input_type = tf.int8
    # 设置推理输出类型为int8
    converter.inference_output_type = tf.int8
    # 执行转换
    tflite_quant_model = converter.convert()

    # 将转换后的模型写入文件
    open(tflite_file, 'wb').write(tflite_quant_model)


# 将教师模型转换为TFLite格式并保存
convert_to_tflite(teacher_model, "teacher.tflite")
# 将学生模型转换为TFLite格式并保存
convert_to_tflite(student_model, "student.tflite")

# 列出当前目录下所有的tflite文件及其详细信息
os.system("ls -lh *.tflite")

# 输出tensorflow的版本（在google colab上使用的2.20.0）
print(f"tensorflow's version：{tf.version.VERSION}")


# 绘制训练损失曲线
# 创建一个大小为10x6的画布
plt.figure(figsize=(10, 6))
# 绘制训练集损失曲线
# history.history['loss'] 和 history.history['val_loss'] 都来自 teacher_model.fit(...) 返回的 History 对象，
#   其中前者记录每个 epoch 的训练损失，后者记录每个 epoch 的验证损失。
plt.plot(history.history['loss'], label='Train Loss', color='blue')
# 绘制验证集损失曲线，使用红色虚线以便区分
plt.plot(history.history['val_loss'], label='Validation Loss', color='red', linestyle='--')
# 设置图表标题
plt.title('Training and Validation Loss over Epochs')
# 设置x轴标签为训练轮次
plt.xlabel('Epoch')
# 设置y轴标签为损失值
plt.ylabel('Loss')
# 显示图例
plt.legend()
# 显示网格，便于观察曲线变化趋势
plt.grid(True)
# 展示图表
plt.show()

最后绘制的训练损失曲线图如下：

image-20260531151902519

三、详细解释版文件代码

# -*- coding: utf-8 -*-
"""
此文件代码已经上传保存到google colab上，可以直接运行。
colab地址为：https://colab.research.google.com/drive/1sFxDyFvHRZr5YvV8S2PFw-583RaucKS3
运行的主要环境包括：Python 3.12.13、tensorflow 2.20.0
"""

# 导入必要的库
import os
import tensorflow as tf

from tensorflow.keras import models
from tensorflow.keras import layers

# 设置随机种子以保证结果可复现
"""
因为设置了一个固定的随机数种子666，因为这个种子是固定的，所以可以保证不论我后续运行此程序几遍，产生的结果都是一样的。
再具体点：通过设置一个固定的随机数种子（例如 666），TensorFlow中所有依赖随机性的操作（比如模型权重的初始化、
    数据的随机打乱等）都会在每次运行时生成相同的随机序列。这正是为了确保实验的可复现性，在多次运行程序时都能得到一致的结果。
"""
tf.random.set_seed(666)

# 加载FashionMNIST数据集，并对像素值进行缩放
# 从 Keras/TensorFlow 预设的网址下载 FashionMNIST 的图片文件和标签文件
# 网址：https://storage.googleapis.com/tensorflow/tf-keras-datasets/
(X_train, y_train), (X_test, y_test) = tf.keras.datasets.fashion_mnist.load_data()

"""
为了直观理解 X_train，可以把它想象成一本有 60,000 页的“方格本”。

先补充一句：
    `X_train`、`X_test` 存的是图片本身；
    `y_train`、`y_test` 存的是每张图片对应的“标准答案”。

关于 y_train / y_test，它们长这样：
    - `y_train.shape` 是 `(60000,)`
    - `y_test.shape` 是 `(10000,)`
    - 它们都是一维数组
    - 数组里的每一个数字，表示一张图片所属的类别编号

例如：
    如果 `y_train[0] = 9`，意思不是“像素值是 9”，
    而是“第 0 张训练图片的类别标签编号是 9”。

在 FashionMNIST 中，标签编号一共是 10 类，通常对应：
    0 = T-shirt/top
    1 = Trouser
    2 = Pullover
    3 = Dress
    4 = Coat
    5 = Sandal
    6 = Shirt
    7 = Sneaker
    8 = Bag
    9 = Ankle boot

一句话理解：
    X 是“题目图片”，y 是“正确答案标签”。

1. 归一化之前：原始数据长什么样
    形状：
        (60000, 28, 28)
    含义：
        - 一共有 60,000 张图片。
        - 每张图片是一个 28 x 28 的像素网格。
        - 每个像素是 0 到 255 之间的整数。

    像素值的含义：
        - 0 表示纯黑色背景。
        - 255 表示最亮的白色。
        - 中间值表示不同深浅的灰色。

    例如某个 3 x 3 局部区域可能是：
        [  0,   0,   0]
        [ 50, 200, 255]
        [ 60, 210, 240]
    这时候它本质上是一个“整数矩阵”。

2. 归一化之后：发生了什么
    执行 X_train = X_train / 255. 之后：
        - 形状还是 (60000, 28, 28)，结构完全没变。
        - 变化的是数值范围：从 0~255 变成 0.0~1.0。

    对应关系大致是：
        - 0 -> 0.0
        - 255 -> 1.0
        - 127.5 -> 0.5

    同样的区域会变成：
        [0.00, 0.00, 0.00]
        [0.20, 0.78, 1.00]
        [0.23, 0.82, 0.94]
    这时候它变成了更适合神经网络处理的“浮点数矩阵”。

3. 为什么要做这一步
    这一步叫做“归一化（Normalization）”。
    对人眼来说，0~255 和 0~1 的图片看起来几乎一样；
    但对神经网络来说，小范围浮点数更容易训练，通常会带来：
        - 更稳定的梯度；
        - 更快的收敛速度；
        - 更好的训练效果。

一句话总结：
    这一行代码是在把“人类习惯的像素整数图片”转换成“模型更容易学习的标准输入”。
"""
# 将训练集像素值缩放到0-1之间
X_train = X_train/255.
# 将测试集像素值缩放到0-1之间
X_test = X_test/255.

# 打印训练集和测试集的形状
X_train.shape, X_test.shape, y_train.shape, y_test.shape

# 将像素值转换为float32类型，并重塑输入数据的形状
"""
这两行代码主要做了两件事：
    1. 数据类型转换；
    2. 形状调整（增加通道维度）。

目的很简单：
    把原始图片整理成卷积神经网络（CNN）能直接接收的标准输入格式。

1. astype("float32")：统一数据类型
    含义：
        把数据强制转换成 32 位浮点数。

    这里要特别注意一个细节：
        - 刚从 `load_data()` 读出来时，`X_train` 通常是 `uint8`；
        - 但在前面执行 `X_train = X_train / 255.` 之后，
          因为参与运算的是浮点数 `255.`，结果通常会变成 `float64`；
        - 所以运行到这里时，`astype("float32")` 实际上是把 `float64`
          再转换成更适合深度学习计算的 `float32`。

    这样做的原因：
        - float32 是深度学习里最常见、最通用的格式；
        - 比 float64 更省显存、计算更快；
        - TensorFlow 中很多层默认就更适配 float32。

2. reshape(-1, 28, 28, 1)：补上通道维度
    修改前：
        (60000, 28, 28)
    可以把它理解成“一叠平面照片”。

    修改后：
        (60000, 28, 28, 1)
    可以把它理解成“带厚度的图像块”，最后那个 1 表示通道数。

    为什么一定要加这个 1？
        因为后面要喂给 Conv2D。
        在 TensorFlow / Keras 中，卷积层要求输入格式是：
            (样本数, 高度, 宽度, 通道数)

        - 彩色图像通常是 3 个通道（RGB）；
        - 灰度图像虽然只有一种颜色信息，也必须显式写成 1 个通道。

        如果只保留 (28, 28)，卷积层会因为找不到“通道维”而报错。

3. 这里的 -1 是什么意思
    -1 表示“让程序自动推断样本数”。
    也就是说，我们只明确告诉它每张图片要变成 28 x 28 x 1，
    剩下前面有多少张样本，程序自己算。

一句话总结：
    这一步是在做“输入格式标准化”，把普通图片整理成 CNN 可以直接处理的 4 维张量。
"""
X_train = X_train.astype("float32").reshape(-1, 28, 28, 1)
X_test = X_test.astype("float32").reshape(-1, 28, 28, 1)

# 定义构建基本浅层卷积神经网络的工具函数
"""
这 8 个 model.add(xxx) 语句就像是在搭积木或者布置流水线。
models.Sequential() 创建了一个空的“架子”，而每一个 model.add() 都是往这个架子上按顺序放上一层层的功能模块。每层处理完数据后，会把结果直接传给下一层。
这 8 层在一起构建了一个 卷积神经网络 (CNN)，我们可以把它分为三个阶段来理解：

第一阶段：提取特征（“看图”）
    这部分负责把图片里的线条、形状找出来。
    1. Conv2D(16, ...)：卷积层（眼睛）
        拿着 16 个不同的“放大镜”（卷积核）在图片上扫描。
        目的是发现图像的局部特征（比如边缘、线条）。
    2. MaxPooling2D(...)：池化层（概括）
        把它看到的东西缩小一半（2x2）。
        它的具体做法是：每次取一个 2x2 的小区域，并且对每一个通道分别取这 4 个位置里的最大值。
        也就是说，它不是直接保留某一个空间位置整组卷积结果，而是按通道分别保留“这一小块区域里最显眼的响应”。
        目的是减少数据量，只保留最显著的特征（比如“这里有个角”，不用管具体在哪一个像素点），防止关注过细的细节。
        例如，如果输入是 (24, 24, 16)，那么经过 2x2 最大池化后会变成 (12, 12, 16)：
            - 24 x 24 变成 12 x 12，是因为长和宽都按 2x2 小块做压缩，相当于各缩小一半；
            - 16 保持不变，是因为池化只是在每一张特征图内部做下采样，不会改变特征图的数量。
    3. Conv2D(32, ...)：卷积层（更高级的眼睛）
        这次用 32 个“放大镜”看刚才概括过的图。
        目的是组合之前的特征，识别更复杂的形状（比如圆形、方形、纹理）。
    4. MaxPooling2D(...)：池化层（再次概括）
        再次把图缩小一半。这时候原本 28x28 的大图已经变得很小很厚了（特征图）。
        这一步的输出数据形状为： (4, 4, 32)
第二阶段：过渡与防作弊
    5. Dropout(0.2)：随机丢弃（防作弊/防止过拟合）
        在训练时，随机让 20% 的神经元“休息”不工作。
        目的是强迫模型不要依赖某几个特定的特征，增强鲁棒性（就像不要只靠记某一道题的答案来考试）。
        将 20% 的神经元随机置零，其他神经元保持不变。
    6. Flatten()：压平（整队）
        前面的操作出来的结果是三维的“立方体”（长x宽x通道）。
        这一步把它拍扁成一条长长的一维数组（向量）。为什么要拍扁？因为后面的全连接层只能吃一维的数据。
        这一步的把 32 张 4x4 特征图首尾相接排成一条长向量（即(4, 4, 32)变成(512,) ），不会改变数值本身。
第三阶段：分类决策（“思考与输出”）
    这部分模拟人脑进行最终判断。
    7. Dense(128, ...)：全连接层（大脑思考）
        用了 128 个神经元来综合分析这一长串特征。
        这是进行逻辑推理的核心区域。
    8. Dense(10)：输出层（给出答案）
        因为我们有 10 类衣服（鞋子、外套等），所以这里有 10 个神经元。
        每个神经元输出一个数值，代表属于某一类的“打分”（分越高越可能是这一类）。
总结
    这 8 行代码构建了一个经典的 Conv-Pool-Conv-Pool-Flatten-Dense 结构。 作为一个 “教师模型” (Teacher Model)，它的结构相对完整且有一定的深度，能够学到比较好的特征，从而能在后面“教”那个只有几层简单的简单学生模型。
"""
def get_teacher_model():
    # 初始化顺序模型
    model = models.Sequential()
    # 添加卷积层，16个卷积核，大小5x5，ReLU激活，输入形状(28, 28, 1)
    # 输出形状会变成 (24, 24, 16)：
    # - 24 x 24 是因为默认 strides=1、padding="valid"，所以空间尺寸按 (28 - 5) / 1 + 1 = 24 计算
    # - 16 表示这一层有16个不同的卷积核，因此会得到16张特征图
    """
    stride=1 的起点滑动表:
    起点    覆盖区间
    0       0  ~ 4
    1       1  ~ 5
    2       2  ~ 6
    3       3  ~ 7
    4       4  ~ 8
    5       5  ~ 9
    6       6  ~ 10
    7       7  ~ 11
    8       8  ~ 12
    9       9  ~ 13
    10      10 ~ 14
    11      11 ~ 15
    12      12 ~ 16
    13      13 ~ 17
    14      14 ~ 18
    15      15 ~ 19
    16      16 ~ 20
    17      17 ~ 21
    18      18 ~ 22
    19      19 ~ 23
    20      20 ~ 24
    21      21 ~ 25
    22      22 ~ 26
    23      23 ~ 27
    24      24 ~ 28  -> 越界，不合法
    所以：
    - 合法起点是 0 ~ 23
    - 一共 24 个位置
    - 输出尺寸就是 24
    """
    model.add(layers.Conv2D(16, (5, 5), activation="relu",
        input_shape=(28, 28, 1)))
    # 添加最大池化层，池化窗口2x2
    # 输出形状会从 (24, 24, 16) 变成 (12, 12, 16)：
    # - 12 x 12 是因为池化窗口是 2 x 2，默认 strides=2，相当于长宽各缩小一半
    # - 16 保持不变，因为池化只在每张特征图内部做下采样，不会改变特征图的张数
    model.add(layers.MaxPooling2D(pool_size=(2, 2)))
    # 添加卷积层，32个卷积核，大小5x5，ReLU激活
    # 输入形状是 (12, 12, 16)，输出形状会变成 (8, 8, 32)：
    # - 8 x 8 是因为默认 strides=1、padding="valid"，所以空间尺寸按 (12 - 5) / 1 + 1 = 8 计算
    # - 32 表示这一层有32个卷积核，因此会得到32张特征图
    # - 这里每个卷积核虽然代码里只写了 (5, 5)，但实际会覆盖 16 个输入通道，所以可理解为 5 x 5 x 16
    model.add(layers.Conv2D(32, (5, 5), activation="relu"))
    # 添加最大池化层，池化窗口2x2
    # 输出形状会从 (8, 8, 32) 变成 (4, 4, 32)：
    # - 4 x 4 是因为池化窗口是 2 x 2，默认 strides=None，此时会自动等于 pool_size=(2, 2)，相当于长宽各缩小一半
    # - 32 保持不变，因为池化只在每张特征图内部做下采样，不会改变特征图的张数
    model.add(layers.MaxPooling2D(pool_size=(2, 2)))
    # 添加Dropout层，丢弃率为0.2，防止过拟合
    # 在这一步之前，单个样本的形状是 (4, 4, 32)，也就是 32 张 4x4 特征图，共 4 x 4 x 32 = 512 个激活值
    # Dropout(0.2) 会在训练时按元素随机将其中约 20% 的激活值置零，以减少模型对局部特征的过度依赖
    model.add(layers.Dropout(0.2))
    # 将多维输出展平为一维
    # 这里会把单个样本的形状从 (4, 4, 32) 变成 (512,)
    # 这一步只是把 32 张 4x4 特征图首尾相接排成一条长向量，不会改变数值本身
    model.add(layers.Flatten())
    # 添加全连接层，128个神经元，ReLU激活
    # 这里会把单个样本的形状从 (512,) 变成 (128,)
    # 可以理解为把前面提取到的 512 个特征综合起来，压缩成 128 个更高层次的特征表示
    # 其中 ReLU 是激活函数，它把 128 个输出里的负数变成 0，正数保持不变 ，这一步可以加速模型的训练 
    model.add(layers.Dense(128, activation="relu"))
    # 添加输出层，10个神经元（对应10个类别）
    # 这里会把单个样本的形状从 (128,) 变成 (10,)
    # 输出的 10 个数分别对应 10 个类别的分数，因此这一层就是最终的分类输出层
    model.add(layers.Dense(10))
    
    # 返回构建好的模型
    return model

# 定义损失函数和优化器
# 使用稀疏分类交叉熵损失函数，from_logits=True表示输入是未经softmax处理的logits
loss_func = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
# 使用Adam优化器
optimizer = tf.keras.optimizers.Adam()

# 准备TensorFlow数据集
# 创建训练数据管道：从tensor切片创建dataset，打乱缓冲区大小100，批次大小64
"""
这段代码是在建立数据流水线（Data Pipeline）。
如果把训练模型比作“喂模型吃饭”，那这一步就是把原始食材加工成一盒一盒能直接送上流水线的便当。

这一行代码可以拆成 3 个动作：

1. from_tensor_slices(...)
    作用：
        把大数组拆成一个个独立样本，形成可迭代的数据流。

    可以理解成：
        原来是一整箱图片，现在变成一张一张排队等待取用。

2. shuffle(100)
    作用：
        在训练时打乱样本顺序，避免模型记住数据原本的排列规律。

    这里的 100 是缓冲区大小：
        - 数值越大，打乱得通常越充分；
        - 但越大也越占内存。

3. batch(64)
    作用：
        每次取 64 张图片和对应标签，打成一个 batch 再送进模型。

    为什么要这样做：
        - 一次只送 1 张图，速度太慢；
        - 一次送一批，可以更好利用 GPU / 向量化计算能力。

总结：
    - `from_tensor_slices`：拆分样本流
    - `shuffle(100)`：随机打乱顺序
    - `batch(64)`：按批次喂给模型

最终得到的 train_ds：
    它不是一次性把所有数据都交给模型，而是每次吐出一个 batch，
    也就是 64 张乱序图片及其对应标签，直到整个数据集遍历完毕。
"""
train_ds = tf.data.Dataset.from_tensor_slices((X_train, y_train)).shuffle(100).batch(64)
# 创建测试数据管道：从tensor切片创建dataset，批次大小64
test_ds = tf.data.Dataset.from_tensor_slices((X_test, y_test)).batch(64)

# 训练教师模型
# 获取教师模型实例
teacher_model = get_teacher_model()
# 编译模型，指定损失函数、优化器和评估指标
"""
这一段需要理解两个点：
    1. `test_ds` 在这里被一份数据复用了两次；
    2. `metrics=["accuracy"]` 到底是怎么计算的。

1. test_ds 在这里扮演了两个角色
    在 `teacher_model.fit(..., validation_data=test_ds)` 中：
        它充当“验证集（validation set）”。
        也就是说，每个 epoch 训练完后，模型都会在这份数据上测一次，
        用来观察当前 loss 和 accuracy 的变化。

    在后面的 `teacher_model.evaluate(test_ds)` 中：
        它又被当成“测试集（test set）”来给出最终成绩。

    补充说明：
        在严格的机器学习实践中，通常会拆成 3 份数据：
        - 训练集：负责学习参数；
        - 验证集：负责训练过程中的监控和调参；
        - 测试集：只在最后做一次客观评估。

    这个脚本为了演示方便，直接让 test_ds 同时兼任验证和测试。

2. metrics=["accuracy"] 是怎么工作的
    当 compile 中写 `metrics=["accuracy"]` 时，
    Keras 会根据任务类型自动选择合适的准确率指标。

    因为这里的损失函数是 `SparseCategoricalCrossentropy`，
    所以 Keras 实际上会使用 `SparseCategoricalAccuracy`。

    它的计算逻辑是：
        1. 对每个样本的输出 logits 取 argmax，得到预测类别索引；
        2. 与真实标签索引做比较；
        3. 预测正确记 1，错误记 0；
        4. 最后对一个 batch 或一个 epoch 内的结果求平均。

3. 为什么训练过程中也能看到 accuracy
    这很好理解成“随堂测验”。

    模型不是一口气学完整个训练集，而是按 batch 一批一批地学：
        - 先对当前 batch 做预测；
        - 立刻就能把预测结果和标签进行比较；
        - 因此也就能立刻算出这一批的准确率。

    然后模型再根据 loss 反向传播更新参数。
    等整个 epoch 的所有 batch 都跑完以后，
    再把这些批次上的表现汇总平均，就形成日志里的训练 accuracy。

一句话总结：
    - `accuracy` 更像训练过程中的“随堂测验平均分”；
    - `val_accuracy` 更像每个 epoch 结束后的“阶段小测”。
"""
teacher_model.compile(loss=loss_func, optimizer=optimizer, metrics=["accuracy"])

# 训练模型，使用 train_ds 作为训练数据， test_ds 作为验证数据，训练10个轮次
history = teacher_model.fit(train_ds,
                  validation_data=test_ds,
                  epochs=10)

# 评估模型并保存权重
# 打印在测试集上的准确率，保留两位小数
# 其中 evaluate(test_ds)[1] 取的是评估结果里的 accuracy，再乘以 100 转成百分数显示
print("Test accuracy: {:.2f}".format(teacher_model.evaluate(test_ds)[1]*100))
# 保存教师模型权重到文件
teacher_model.save_weights("teacher_model.weights.h5")

# 学生模型工具函数
"""
这里把学生模型设计得很简单，是故意的，不是偷懒。

知识蒸馏的核心思想是：
    “让一个强模型把自己的经验传给一个弱模型”。

为什么学生模型要故意做得很弱？

1. 为了体现“模型压缩”的价值
    教师模型使用卷积层、池化层，参数更多、表达能力更强，
    但代价是计算更慢、内存占用更大。

    学生模型则故意走“极简路线”：
        - 不用卷积层；
        - 直接把图片拍平成向量；
        - 只保留很小的全连接层。

    这样做可以模拟轻量设备场景，比如手机端、嵌入式设备、低功耗设备等。

2. 为了更明显地展示蒸馏效果
    如果学生模型本身也很强，那它即使不靠老师也可能学得不错，
    就很难看出“蒸馏到底帮了多少”。

    现在反过来：
        学生模型天生比较弱，如果它在老师指导下依然能取得不错效果，
        就能更有说服力地说明蒸馏确实传递了有价值的“暗知识”。

3. 这其实是一种“极限测试”
    这个 Demo 最有意思的地方就在于结构反差非常大：
        - Teacher：会“看图”的卷积网络；
        - Student：几乎退化成最基础的全连接网络。

    如果这样一个非常轻量、先天不擅长处理图像的学生模型，
    仍然能在教师指导下学会分类，
    那就说明蒸馏不仅传递了“标准答案”，
    还传递了类别之间更细腻的相似性信息。

一句话总结：
    学生模型之所以故意设计得简单，是为了同时展示两件事：
        - 蒸馏可以帮助模型压缩；
        - 蒸馏可以让弱模型学到强模型的经验。
"""
def get_student_model():
    # 初始化顺序模型
    model = models.Sequential()
    # 添加输入层，形状为(28, 28, 1)
    model.add(layers.Input(shape=(28, 28, 1)))
    # 将输入展平
    model.add(layers.Flatten())
    # 添加全连接层，48个神经元，ReLU激活
    model.add(layers.Dense(48, activation="relu"))
    # 添加输出层，10个神经元
    model.add(layers.Dense(10))

    # 返回构建好的模型
    return model

# 来源：https://github.com/google-research/simclr/blob/master/colabs/distillation_self_training.ipynb
def get_kd_loss(student_logits, teacher_logits, temperature=0.5):
    """
    知识蒸馏的核心损失函数。
    作用：计算学生模型和教师模型在“软目标（Soft Targets）”上的差异。

    参数解析：
    1. temperature (温度 T): 这是一个超参数，用于控制概率分布的“平滑程度”。
       - 当 T=1 时：标准的 Softmax。概率分布非常尖锐（Teacher非常自信，通过Softmax后接近One-hot编码）。
         例如：[猫:0.99, 狗:0.009, 汽车:0.001]。
         这时候学生只能学到“它是猫”，学不到别的。
       - 当 T 较大（如5.0）时：概率分布变平滑（Teacher变得“犹豫”）。
         例如：[猫:0.6, 狗:0.3, 汽车:0.1]。
         Teacher 透露了“暗知识”：虽然它是猫，但它长得挺像狗，完全不像汽车。
         学生模仿这个分布，就能学到“猫和狗具有相似性”这种深层结构知识。

    逻辑步骤：
    1. teacher_probs: 将 Teacher 的 Logits 除以 T，然后 Softmax。得到“软标签”。
    2. student_logits / temperature: 学生的 Logits 也必须除以同样的 T，以便在同一尺度上对比。
    3. loss: 使用交叉熵衡量两个分布的差异。
    4. temperature**2: 梯度补偿。
       - 因为 Logits 被除以了 T，导致计算出的梯度值会缩小 T^2 倍。
       - 为了防止梯度太小导致学生学不动，我们需要把 Loss 乘以 T^2 补回来。

    数值示例 (假设 T=4):
    -------------
    Step 1: 准备原始数据 (Logits)
      Teacher: [狗: 10.0,  猫: 6.0,   汽车: -4.0]  <-- 教师非常自信，10.0远大于其他
        #import torch
        #scores = torch.tensor([10.0, 6.0, -4.0])
        #probs = torch.nn.functional.softmax(scores, dim=0)
        #print(probs)  # tensor([9.8201e-01, 1.7986e-02, 8.1444e-07]) #就是说，如果不进行温度缩放，猫的概率几乎为0
      Student: [狗: 4.0,   猫: 2.0,   汽车: -2.0]  <-- 学生还在学习中
      * 注意：如果此时直接用 T=1 做 Softmax:
        Teacher_prob(猫) ≈ 0.018 (1.8%)。教师几乎完全忽略了猫的可能性，学生学不到太多细节。

    Step 2: 温度缩放 (Logits / T)
      动作：将所有 Logits 除以 4。
      Teacher: [2.5,  1.5,  -1.0]  <-- 分数差距被拉近了 (原本差4分，现在只差1分)
      Student: [1.0,  0.5,  -0.5]

    Step 3: 软化概率 (Softmax)
      动作：对缩放后的 Logits 进行 Softmax。
      Teacher_probs: [0.71, 0.26, 0.02] 
      * 重点：现在老师指出“虽是狗，但有 26% 的概率像猫”。这就是被放大的“暗知识”。
      Student_probs: [0.53, 0.32, 0.15]

    Step 4: 计算损失 (Cross Entropy)
      目的：衡量“学生的概率分布”是否对齐了“教师的概率分布”。
      公式：Loss = -sum(Teacher_probs * log(Student_probs))
      
      * 这里的 log 是“逐元素”对所有概率求自然对数：
        log([0.53, 0.32, 0.15]) -> [-0.63, -1.14, -1.90]
      * 然后对应位置相乘（Teacher * log_Student）：
        - 狗: 0.71 * -0.63 = -0.45
        - 猫: 0.26 * -1.14 = -0.30
        - 车: 0.02 * -1.90 = -0.04
      * 最后求和并取反：
        Sum = -(-0.45 - 0.30 - 0.04) = -(-0.79) = 0.79 -> Loss = 0.79

      * 老师觉得猫有 0.26，学生预测有 0.32，两者不一致 -> 产生 Loss。
      * 优化这个 Loss 能让学生去模仿老师的这种“犹豫感”。

    Step 5: 梯度补偿 (Loss * T^2)
      动作：最终 Loss 乘以 16 (因为 T=4, T^2=16)。
      原因：Step 2 中把 Logits 除以了 4，这导致反向传播回来的梯度自然地缩小了约 16 倍。
             为了防止梯度太小导致学生学不动（参数更新太慢），必须人为地乘回来，保持梯度的量级正常。
    """
    # 计算教师模型的软标签概率，使用温度参数缩放logits
    teacher_probs = tf.nn.softmax(teacher_logits / temperature)
    # 计算知识蒸馏损失（使用兼容v1版本的loss函数，计算softmax交叉熵）
    # 注意这里student_logits也除了温度，并且损失乘以温度的平方以保持梯度量级
    kd_loss = tf.compat.v1.losses.softmax_cross_entropy(teacher_probs, student_logits / temperature, temperature**2)
    return kd_loss

# 模型和优化器
# 获取学生模型实例
student_model = get_student_model()
# 定义优化器，学习率为0.01
optimizer = tf.keras.optimizers.Adam(learning_rate=0.01)

# Average the loss across the batch size within an epoch
# 计算每个epoch内batch的平均损失
train_loss = tf.keras.metrics.Mean(name="train_loss")
valid_loss = tf.keras.metrics.Mean(name="test_loss")

# 指定性能指标
"""
SparseCategoricalAccuracy 的计算逻辑：
    它专门用于比较“整数标签”和“模型输出”是否匹配。

    具体步骤：
        1. labels 是真实类别索引，例如 [5, 2]；
        2. predictions 是模型输出的 logits 或概率；
        3. 对 predictions 做 argmax，取最大值所在类别作为预测结果；
        4. 将预测类别和真实标签逐个对比；
        5. 统计整体平均正确率。

    例如：
        labels = [5, 2, 7]
        argmax(predictions) = [5, 2, 9]
        其中前两张预测正确，第三张错误，所以准确率是 2 / 3。
"""
train_acc = tf.keras.metrics.SparseCategoricalAccuracy(name="train_acc")
valid_acc = tf.keras.metrics.SparseCategoricalAccuracy(name="valid_acc")

# 训练工具函数
@tf.function
def model_train(images, labels, teacher_model,
                student_model, optimizer, temperature):
    # 获取教师模型对图像的预测logits（不进行梯度追踪）
    teacher_logits = teacher_model(images)

    # 开启梯度记录
    with tf.GradientTape() as tape:
        # 获取学生模型对图像的预测logits
        student_logits = student_model(images)
        # 计算知识蒸馏损失
        loss = get_kd_loss(student_logits, teacher_logits, temperature)

        # 计算相对于学生模型可训练变量的梯度
        # 这里的意思是：以 loss 为目标，分别计算它对学生模型每个可训练参数的导数。
        # 这些导数就叫梯度，可以理解为“这个参数往哪个方向改、改多大，能让 loss 变小”。
        # student_model.trainable_variables 里包含的是学生模型中所有要学习的参数，例如卷积核权重、全连接层权重和偏置。
        # tape.gradient(...) 执行完后，gradients 会得到一个列表，里面每一项都和某个可训练参数一一对应。
        #
        # 一个极小例子：
        #   假设某个参数当前是 w = 0.8
        #   计算得到它对应的梯度是 dw = 0.2
        #   这可以粗略理解为：如果想让 loss 下降，w 往更小的方向调整会更合适
        gradients = tape.gradient(loss, student_model.trainable_variables)
        # 应用梯度更新学生模型的参数
        # 这里会把 gradients 和 student_model.trainable_variables 一一配对后交给优化器。
        # zip(...) 可以理解为得到 [(梯度1, 参数1), (梯度2, 参数2), ...] 这样的对应关系。
        # 然后 optimizer 会根据自己的更新规则去修改这些参数；如果粗略理解成最基础的梯度下降，就是：
        #   新参数 = 旧参数 - 学习率 x 梯度
        #
        # 继续上面的例子：
        #   假设学习率 lr = 0.1
        #   那么更新后大致变成：
        #   w_new = 0.8 - 0.1 x 0.2 = 0.78
        # 也就是说，这两行代码合起来做的事就是：
        #   1. 先算出每个参数该怎么改
        #   2. 再真的按这个方向去更新参数
        optimizer.apply_gradients(zip(gradients, student_model.trainable_variables))

        # 更新训练损失指标
        # train_loss(loss) 就是把当前 batch 的损失值交给 train_loss 这个平均值统计器，让它持续累计，后面就能得到整个训练过程的平均损失。
        train_loss(loss)
        # 更新训练准确率指标
        # 把当前 batch 的预测结果和真实标签记到“训练准确率统计器”里，让它持续累计，后面就能得到整个训练过程的平均准确率。
        #   train_acc 需要两个参数 因为准确率不像损失那样已经提前算成一个数了，它需要现场比较：真实标签 labels 与 模型预测 tf.nn.softmax(student_logits)
        #   然后内部会做类似这样的事：（1）对预测结果取 argmax （2）和 labels 逐个比较 （3）统计这一批里预测正确的比例 （4）再累计到整个 epoch 的准确率统计中
        train_acc(labels, tf.nn.softmax(student_logits))

# 验证工具函数
@tf.function
def model_validate(images, labels, teacher_model,
                   student_model, temperature):
    # 获取教师模型预测logits
    teacher_logits = teacher_model(images)

    # 获取学生模型预测logits
    student_logits = student_model(images)
    # 计算知识蒸馏损失
    loss = get_kd_loss(student_logits, teacher_logits, temperature)

    # 更新验证损失指标
    valid_loss(loss)
    # 更新验证准确率指标
    valid_acc(labels, tf.nn.softmax(student_logits))

# 将所有步骤结合在一起进行训练
def train_model(epochs, teacher_model, student_model, optimizer, temperature=0.5):
    # 遍历每个epoch
    for epoch in range(epochs):
        # 遍历训练数据集中的每个batch
        # 这里的 images, labels 长度通常为 64 (因为我们在定义 train_ds 时设置了 .batch(64))
        # 唯一的例外是最后一个 batch，如果数据总数无法整除 64，剩下的余数会在最后一次取出。
        for (images, labels) in train_ds:
            # 执行单步训练
            model_train(images, labels, teacher_model, student_model, optimizer, temperature)

        # 遍历测试数据集中的每个batch
        for (images, labels) in test_ds:
            # 执行单步验证
            model_validate(images, labels, teacher_model, student_model, temperature)

        # 获取当前epoch的平均损失和准确率
        (loss, acc) = train_loss.result(), train_acc.result()
        (val_loss, val_acc) = valid_loss.result(), valid_acc.result()

        # 重置指标状态，为下一个epoch做准备
        train_loss.reset_state(), train_acc.reset_state()
        valid_loss.reset_state(), valid_acc.reset_state()

        # 定义打印模板
        template = "Epoch {}, loss: {:.3f}, acc: {:.3f}, val_loss: {:.3f}, val_acc: {:.3f}"
        # 打印当前epoch的训练结果
        print (template.format(epoch+1,
                            loss,
                            acc,
                            val_loss,
                            val_acc))


    # 返回训练好的教师和学生模型
    return teacher_model, student_model

# 开始训练模型，训练10个epoch
_, student_model = train_model(10, teacher_model, student_model, optimizer)

# 序列化保存
# 保存学生模型权重
student_model.save_weights("student_model.weights.h5")

# Investigate the sizes
# 查看模型文件大小
# 列出当前目录下所有的h5文件及其详细信息
os.system("ls -lh *.h5")

# 打印教师模型摘要
# 通过摘要查看教师模型的层结构、输出形状和参数数量
teacher_model.summary()

# 打印学生模型摘要
# 通过摘要查看学生模型的层结构、输出形状和参数数量
student_model.summary()

"""
使用TFLite可以进一步减小模型大小。
"""
# 来源：https://www.tensorflow.org/lite/performance/post_training_quant

# 定义代表性数据集生成器函数
def representative_data_gen():
    # 从训练数据中取100个样本，每次一个，用于量化校准
    for input_value in tf.data.Dataset.from_tensor_slices(X_train).batch(1).take(100):
        yield [input_value]

# 定义转换为TFLite模型的函数
def convert_to_tflite(model, tflite_file):
    # 从Keras模型创建TFLite转换器
    converter = tf.lite.TFLiteConverter.from_keras_model(model)
    # 设置优化选项为默认优化
    converter.optimizations = [tf.lite.Optimize.DEFAULT]
    # 设置代表性数据集用于量化
    converter.representative_dataset = representative_data_gen
    # 设置目标规范支持的操作集为INT8内置操作
    converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
    # 设置推理输入类型为int8
    converter.inference_input_type = tf.int8
    # 设置推理输出类型为int8
    converter.inference_output_type = tf.int8
    # 执行转换
    tflite_quant_model = converter.convert()

    # 将转换后的模型写入文件
    open(tflite_file, 'wb').write(tflite_quant_model)

# 将教师模型转换为TFLite格式并保存
convert_to_tflite(teacher_model, "teacher.tflite")
# 将学生模型转换为TFLite格式并保存
convert_to_tflite(student_model, "student.tflite")

# 列出当前目录下所有的tflite文件及其详细信息
os.system("ls -lh *.tflite")

# 输出tensorflow的版本（在google colab上使用的2.20.0，以及Python 3.12.13）
print(f"tensorflow's version：{tf.version.VERSION}")

# 绘制训练损失曲线
# 创建一个大小为10x6的画布
plt.figure(figsize=(10, 6))
# 绘制训练集损失曲线
# history.history['loss'] 和 history.history['val_loss'] 都来自 teacher_model.fit(...) 返回的 History 对象，
#   其中前者记录每个 epoch 的训练损失，后者记录每个 epoch 的验证损失。
plt.plot(history.history['loss'], label='Train Loss', color='blue')
# 绘制验证集损失曲线，使用红色虚线以便区分
plt.plot(history.history['val_loss'], label='Validation Loss', color='red', linestyle='--')
# 设置图表标题
plt.title('Training and Validation Loss over Epochs')
# 设置x轴标签为训练轮次
plt.xlabel('Epoch')
# 设置y轴标签为损失值
plt.ylabel('Loss')
# 显示图例
plt.legend()
# 显示网格，便于观察曲线变化趋势
plt.grid(True)
# 展示图表
plt.show()