最近爆火的AI大模型“蒸馏”技术是什么？3分钟带你轻松入门！AI大模型零基础入门到精通实战教程！

如果你最近关注了DeepSeek的最新动态，可能会经常听到“蒸馏”（Distillation）这个词。但它到底是什么意思？为什么它如此重要？在本文中，将解释蒸馏的过程，并通过一个TensorFlow示例进行演示。通过本文的学习，你将更深入地理解和欣赏模型蒸馏的价值。# 教师模型keras.layers.Dense(10) # 不使用softmax，输出原始logits以用于蒸馏])注意，学生模型的

编程瞬息全宇宙

895人浏览 · 2025-02-27 16:43:58

编程瞬息全宇宙 · 2025-02-27 16:43:58 发布

模型蒸馏简介

如果你最近关注了DeepSeek的最新动态，可能会经常听到“蒸馏”（Distillation）这个词。但它到底是什么意思？为什么它如此重要？在本文中，将解释蒸馏的过程，并通过一个TensorFlow示例进行演示。通过本文的学习，你将更深入地理解和欣赏模型蒸馏的价值。

模型蒸馏的工作原理

模型蒸馏是一种技术，通过让一个更小、更简单的模型（学生模型）学习一个更大、更复杂的模型（教师模型）的软化概率输出，而不是仅仅学习原始标签，从而使学生模型能够以更紧凑的形式捕获教师模型的知识。这样，学生模型可以用更少的参数实现与教师模型相似的性能。例如，在图像分类任务中，学生模型不仅学习图像是“狗”还是“猫”，还会从教师模型的置信度分数中学习（例如，80%是狗，15%是猫，5%是狐狸），从而保留更细致的知识。这一过程减少了模型的大小和计算需求，同时保持了较高的准确性。

让我们通过一个示例来了解具体如何实现。我们将使用MNIST数据集训练一个卷积神经网络（CNN）。

MNIST数据集

MNIST数据集（Modified National Institute of Standards and Technology）是机器学习和计算机视觉领域广泛使用的基准数据集。它包含70,000张28x28像素的手写数字灰度图像（0-9），其中60,000张用于训练，10,000张用于测试。

首先，我们有一个教师模型：

在这里插入图片描述
教师模型

教师模型是一个使用MNIST数据集训练的CNN。

我们还有一个学生模型，它是一个比教师模型更简单、更小的模型：

在这里插入图片描述

学生模型

模型蒸馏的目标是训练一个更小的学生模型，使其能够模仿教师模型的性能，同时减少计算和训练时间。

我给大家准备了一份全套的《AI大模型零基础入门+进阶学习资源包》，包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。😝有需要的小伙伴，可以VＸ扫描下方二维码免费领取🆓

接下来，教师模型和学生模型都将使用数据集进行预测。然后，计算它们输出之间的Kullback-Leibler（KL）散度（稍后解释）。该值用于确定梯度——指示模型的每个部分应如何调整——从而使学生模型能够相应地更新：

蒸馏过程

学生模型现在已训练完成，并达到了与教师模型相当的准确性：

在这里插入图片描述
训练结果

创建一个模型蒸馏的示例项目

现在你已经对模型蒸馏的工作原理有了更清晰的理解，接下来我们将通过一个简单的示例来展示如何实现它。为此，我将使用TensorFlow和MNIST数据集来训练一个教师模型，然后应用模型蒸馏技术来训练一个更小的学生模型，使其能够模仿教师模型的性能，同时减少资源需求。

使用MNIST数据集

首先，确保已安装TensorFlow：

!pip install tensorflow

接下来，加载MNIST数据集：

from tensorflow import keras
import matplotlib.pyplot as plt

# 加载数据集（MNIST）
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

fig = plt.figure()

# 可视化一些数字
for i in range(9):
    plt.subplot(3,3,i+1)
    plt.tight_layout()
    plt.imshow(x_train[i], interpolation='none')
    plt.title("Digit: {}".format(y_train[i]))

    # 不显示x和y轴刻度
    plt.xticks([])
    plt.yticks([])

以下是MNIST数据集中的前9个样本数字及其标签：

MNIST样本

你还需要对图像数据进行归一化，并扩展数据集的维度以准备训练：

import tensorflow as tf
import numpy as np

# 归一化图像
x_train, x_test = x_train / 255.0, x_test / 255.0

# 扩展维度以适应CNN
x_train = np.expand_dims(x_train, axis=-1)
x_test = np.expand_dims(x_test, axis=-1)

# 将标签转换为分类格式（one-hot编码）
y_train = keras.utils.to_categorical(y_train, 10)
y_test = keras.utils.to_categorical(y_test, 10)

定义教师模型

接下来，我们定义一个教师模型——一个包含多个层的CNN：

# 教师模型
teacher_model = keras.Sequential([
    keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
    keras.layers.MaxPooling2D((2, 2)),
    keras.layers.Conv2D(64, (3, 3), activation='relu'),
    keras.layers.MaxPooling2D((2, 2)),
    keras.layers.Flatten(),
    keras.layers.Dense(128, activation='relu'),
    keras.layers.Dense(10)  # 不使用softmax，输出原始logits以用于蒸馏
])

注意，学生模型的最后一层有10个单元（每个数字一个），但没有使用softmax激活函数。它输出原始logits，这在模型蒸馏过程中非常重要，因为softmax会在蒸馏过程中应用于计算教师模型和学生模型之间的Kullback-Leibler（KL）散度。

定义好教师神经网络后，使用compile()方法配置优化器、损失函数和评估指标：

teacher_model.compile(
    optimizer = 'adam',
    loss = tf.keras.losses.CategoricalCrossentropy(from_logits = True),
    metrics = ['accuracy']
)

接下来，使用fit()方法训练模型：

# 训练教师模型
teacher_model.fit(x_train, y_train, 
                  epochs = 5, 
                  batch_size = 64, 
                  validation_data = (x_test, y_test))

本次训练使用了5个epoch：

Epoch 1/5
938/938 ━━━━━━━━━━━━━━━━━━━━ 8s 8ms/step - accuracy: 0.8849 - loss: 0.3798 - val_accuracy: 0.9844 - val_loss: 0.0504
Epoch 2/5
938/938 ━━━━━━━━━━━━━━━━━━━━ 9s 9ms/step - accuracy: 0.9847 - loss: 0.0494 - val_accuracy: 0.9878 - val_loss: 0.0361
Epoch 3/5
938/938 ━━━━━━━━━━━━━━━━━━━━ 9s 10ms/step - accuracy: 0.9907 - loss: 0.0302 - val_accuracy: 0.9898 - val_loss: 0.0316
Epoch 4/5
938/938 ━━━━━━━━━━━━━━━━━━━━ 10s 10ms/step - accuracy: 0.9928 - loss: 0.0223 - val_accuracy: 0.9895 - val_loss: 0.0303
Epoch 5/5
938/938 ━━━━━━━━━━━━━━━━━━━━ 10s 11ms/step - accuracy: 0.9935 - loss: 0.0197 - val_accuracy: 0.9919 - val_loss: 0.0230

定义学生模型

教师模型训练完成后，接下来定义学生模型。学生模型是一个比教师模型更简单的架构，层数更少：

# 学生模型
student_model = keras.Sequential([
    keras.layers.Conv2D(16, (3, 3), activation='relu', input_shape=(28, 28, 1)),
    keras.layers.MaxPooling2D((2, 2)),
    keras.layers.Flatten(),
    keras.layers.Dense(64, activation='relu'),
    keras.layers.Dense(10)  # 不使用softmax，输出原始logits以用于蒸馏
])

定义蒸馏损失函数

接下来，定义蒸馏函数，该函数将使用教师模型的预测和学生模型的预测来计算蒸馏损失。该函数将：

使用教师模型生成输入批次的软目标（概率）。
使用学生模型的预测计算其软概率。
计算教师模型和学生模型软概率之间的Kullback-Leibler（KL）散度。
返回蒸馏损失。

软概率指的是多个可能结果的概率分布，而不是分配一个硬标签。例如，在一个预测电子邮件是否为垃圾邮件的机器学习模型中，模型可能会输出一个概率，如0.85表示垃圾邮件，0.15表示非垃圾邮件。这意味着模型有85%的置信度认为邮件是垃圾邮件，但仍认为有15%的可能性不是，从而允许更好的决策和阈值调整。

软概率使用softmax函数计算，并通过温度参数控制。在知识蒸馏中，教师模型的软概率帮助学生模型学习关于类别之间关系的更丰富信息，从而提高泛化能力和性能。

以下是distillation_loss()函数的定义：

def distillation_loss(y_true, y_pred, x_batch, teacher_model, temperature=5):
    """
    使用KL散度计算蒸馏损失。
    """
    # 计算当前批次的教师logits
    teacher_logits = teacher_model(x_batch, training=False)

    # 将logits转换为软概率
    teacher_probs = tf.nn.softmax(teacher_logits / temperature)
    student_probs = tf.nn.softmax(y_pred / temperature)

    # KL散度损失（教师和学生分布之间的差异）
    return tf.reduce_mean(tf.keras.losses.KLDivergence()(teacher_probs, student_probs))

Kullback-Leibler（KL）散度，也称为相对熵，是衡量一个概率分布与另一个参考概率分布之间差异的指标。

使用知识蒸馏训练学生模型

现在，你可以使用知识蒸馏来训练学生模型。首先，定义train_step()函数：

optimizer = tf.keras.optimizers.Adam()

@tf.function
def train_step(x_batch, y_batch, student_model, teacher_model):
    with tf.GradientTape() as tape:
        # 获取学生模型的预测
        student_preds = student_model(x_batch, training=True)

        # 计算蒸馏损失（显式传递teacher_model）
        loss = distillation_loss(y_batch, student_preds, x_batch, teacher_model, temperature=5)

    # 计算梯度
    gradients = tape.gradient(loss, student_model.trainable_variables)

    # 应用梯度——训练学生模型
    optimizer.apply_gradients(zip(gradients, student_model.trainable_variables))

    return loss

该函数执行单次训练步骤：

计算学生模型的预测。
使用教师模型的预测计算蒸馏损失。
计算梯度并更新学生模型的权重。

为了训练学生模型，你创建一个训练循环来迭代数据集，在每一步更新学生模型的权重，并在每个epoch结束时打印损失以监控进度：

# 训练循环
epochs = 5
batch_size = 32

# 准备数据集批次
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(batch_size)

for epoch in range(epochs):
    total_loss = 0
    num_batches = 0

    for x_batch, y_batch in train_dataset:
        loss = train_step(x_batch, y_batch, student_model, teacher_model)
        total_loss += loss.numpy()
        num_batches += 1

    avg_loss = total_loss / num_batches
    print(f"Epoch {epoch+1}, Loss: {avg_loss:.4f}")

print("学生模型训练完成！")

训练完成后，你应该会看到类似以下的结果：

Epoch 1, Loss: 0.1991
Epoch 2, Loss: 0.0588
Epoch 3, Loss: 0.0391
Epoch 4, Loss: 0.0274
Epoch 5, Loss: 0.0236
学生模型训练完成！

评估学生模型

学生模型训练完成后，你可以使用测试集（x_test和y_test）评估其性能：

student_model.compile(
    optimizer='adam',
    loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True),
    metrics=['accuracy']
)

student_acc = student_model.evaluate(x_test, y_test, verbose=0)[1]
print(f"学生模型准确率: {student_acc:.4f}")

正如预期的那样，学生模型达到了相当不错的准确率：

学生模型准确率: 0.9863

使用教师模型和学生模型进行预测

现在，你可以使用教师模型和学生模型进行一些预测，看看两者是否能够准确预测MNIST测试数据集中的数字：

import numpy as np

_, (x_test, y_test) = keras.datasets.mnist.load_data()

for index in range(5):    
    plt.figure(figsize=(2, 2))
    plt.imshow(x_test[index], interpolation='none')
    plt.title("Digit: {}".format(y_test[index]))
    # 不显示x和y轴刻度
    plt.xticks([])
    plt.yticks([])
    plt.show()
    
    # 现在可以进行预测
    x = x_test[index].reshape(1,28,28,1)
    
    predictions = teacher_model.predict(x)
    print(predictions)
    print("教师模型预测值: ", np.argmax(predictions, axis=-1))
    
    predictions = student_model.predict(x)
    print(predictions)
    print("学生模型预测值: ", np.argmax(predictions, axis=-1))

你会发现学生模型的表现与教师模型一样好。