PyTorch深度学习框架60天进阶学习计划 - 第47天：模型压缩蒸馏技术（一）

PyTorch深度学习框架60天进阶学习计划 - 第47天：模型压缩蒸馏技术（一）

第一部分：知识蒸馏的温度调节机制详解

欢迎来到我们学习计划的第47天！今天我们将深入探讨模型压缩技术中的两个重要方法：知识蒸馏和模型剪枝。

在第一部分，我们将聚焦于知识蒸馏的温度调节机制。

1. 知识蒸馏概述

知识蒸馏（Knowledge Distillation）是Geoffrey Hinton在2015年提出的一种模型压缩方法，核心思想是将一个复杂的"教师模型"(Teacher Model)的"知识"迁移到一个更小的"学生模型"(Student Model)中。这里的"知识"不仅仅是最终的预测结果，还包括模型对不同类别的置信度分布。

想象一下，如果我们给教师模型一张猫的图片，它可能会以90%的概率预测为"猫"，5%的概率预测为"狗"，以及其他类别的小概率。这种概率分布包含了教师模型学到的关于输入样本的丰富信息，比如"这张猫的图片有些特征和狗相似"。如果我们只关注硬标签（即预测为概率最高的类别），这些细微但有价值的信息就会丢失。

知识蒸馏的巧妙之处在于，它使学生模型不仅学习真实标签，还学习教师模型的"软标签"（即概率分布），从而继承教师模型的泛化能力。

2. 温度调节机制原理

在知识蒸馏中，“温度”(Temperature)是一个关键参数，用于控制概率分布的"软硬程度"。温度参数T的作用体现在softmax函数上：

softmax(z_i, T) = exp(z_i/T) / sum(exp(z_j/T))

其中，z_i是模型的logits（即最后一层的输出，尚未经过softmax转换），T是温度参数。

温度参数效果分析

• 当T=1时，就是标准的softmax函数
• 当T→0时，分布变得更"硬"，接近one-hot编码（赢家通吃）
• 当T→∞时，分布变得更"软"，接近于均匀分布

通过提高温度T，我们可以"软化"概率分布，使得较小的概率值变得更加明显，这有助于学生模型学习到教师模型中的细微差别。

让我们通过一个具体例子来说明温度参数的影响：

Logits	T=1	T=2	T=10
[10, 5, 1]	[0.952, 0.047, 0.001]	[0.881, 0.112, 0.007]	[0.550, 0.301, 0.149]
[10, 8, 6]	[0.665, 0.244, 0.090]	[0.527, 0.320, 0.153]	[0.375, 0.345, 0.280]

从上表可以看出，随着温度的增加，原本差异很大的概率值变得更接近了。例如，在第一行，当T=1时，第一个类别的概率远高于其他类别；但当T=10时，概率分布变得更加均匀，使得"暗示性"知识更加明显。

3. 知识蒸馏的损失函数

知识蒸馏的损失函数通常由两部分组成：

1. 蒸馏损失 (Distillation Loss)：学生模型输出与教师模型软标签之间的差距
2. 学生损失 (Student Loss)：学生模型输出与真实标签之间的差距

总损失函数为：

L_total = α * L_distill + (1 - α) * L_student

其中，α是一个平衡两种损失的超参数，L_distill通常使用KL散度或交叉熵来计算软标签之间的差距，L_student使用标准的分类交叉熵损失。

4. PyTorch实现知识蒸馏

下面是一个使用PyTorch实现知识蒸馏的完整示例，包括温度调节机制：

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt# 设置随机种子以确保结果可复现
torch.manual_seed(42)# 定义教师模型（较复杂的CNN）
class TeacherModel(nn.Module):def __init__(self):super(TeacherModel, self).__init__()self.conv1 = nn.Conv2d(1, 32, 3, 1)self.conv2 = nn.Conv2d(32, 64, 3, 1)self.dropout1 = nn.Dropout2d(0.25)self.dropout2 = nn.Dropout2d(0.5)self.fc1 = nn.Linear(9216, 128)self.fc2 = nn.Linear(128, 10)def forward(self, x):x = self.conv1(x)x = F.relu(x)x = self.conv2(x)x = F.relu(x)x = F.max_pool2d(x, 2)x = self.dropout1(x)x = torch.flatten(x, 1)x = self.fc1(x)x = F.relu(x)x = self.dropout2(x)x = self.fc2(x)return x# 定义学生模型（较简单的CNN）
class StudentModel(nn.Module):def __init__(self):super(StudentModel, self).__init__()self.conv1 = nn.Conv2d(1, 16, 3, 1)self.fc1 = nn.Linear(13*13*16, 32)self.fc2 = nn.Linear(32, 10)def forward(self, x):x = self.conv1(x)x = F.relu(x)x = F.max_pool2d(x, 2)x = torch.flatten(x, 1)x = self.fc1(x)x = F.relu(x)x = self.fc2(x)return x# 定义知识蒸馏损失函数
class DistillationLoss(nn.Module):def __init__(self, alpha=0.5, temperature=4.0):super(DistillationLoss, self).__init__()self.alpha = alphaself.temperature = temperaturedef forward(self, student_logits, teacher_logits, labels):# 使用高温度的softmax计算软标签soft_targets = F.softmax(teacher_logits / self.temperature, dim=1)soft_prob = F.log_softmax(student_logits / self.temperature, dim=1)# 计算蒸馏损失（KL散度）distillation_loss = F.kl_div(soft_prob, soft_targets, reduction='batchmean') * (self.temperature ** 2)# 计算学生模型与真实标签的交叉熵损失student_loss = F.cross_entropy(student_logits, labels)# 总损失 = α * 蒸馏损失 + (1-α) * 学生损失total_loss = self.alpha * distillation_loss + (1 - self.alpha) * student_lossreturn total_loss, distillation_loss, student_loss# 加载MNIST数据集
def load_data():transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.1307,), (0.3081,))])train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)test_dataset = datasets.MNIST('./data', train=False, transform=transform)train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)return train_loader, test_loader# 训练教师模型
def train_teacher(model, train_loader, test_loader, epochs=5):device = torch.device("cuda" if torch.cuda.is_available() else "cpu")model.to(device)optimizer = optim.Adam(model.parameters(), lr=0.001)for epoch in range(epochs):model.train()for batch_idx, (data, target) in enumerate(train_loader):data, target = data.to(device), target.to(device)optimizer.zero_grad()output = model(data)loss = F.cross_entropy(output, target)loss.backward()optimizer.step()if batch_idx % 100 == 0:print(f'Teacher Training: Epoch {epoch+1}/{epochs} [{batch_idx*len(data)}/{len(train_loader.dataset)} ({100.*batch_idx/len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}')# 在测试集上评估model.eval()test_loss = 0correct = 0with torch.no_grad():for data, target in test_loader:data, target = data.to(device), target.to(device)output = model(data)test_loss += F.cross_entropy(output, target, reduction='sum').item()pred = output.argmax(dim=1, keepdim=True)correct += pred.eq(target.view_as(pred)).sum().item()test_loss /= len(test_loader.dataset)accuracy = 100. * correct / len(test_loader.dataset)print(f'Teacher Test: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)')return model# 知识蒸馏训练学生模型
def train_student_with_distillation(teacher_model, student_model, train_loader, test_loader, temperature=4.0, alpha=0.5, epochs=5):device = torch.device("cuda" if torch.cuda.is_available() else "cpu")teacher_model.to(device)student_model.to(device)teacher_model.eval()  # 教师模型设为评估模式optimizer = optim.Adam(student_model.parameters(), lr=0.001)distillation_criterion = DistillationLoss(alpha=alpha, temperature=temperature)# 记录训练过程中的损失和准确率history = {'total_loss': [], 'distill_loss': [], 'student_loss': [],'test_loss': [], 'test_accuracy': []}for epoch in range(epochs):student_model.train()epoch_total_loss = 0epoch_distill_loss = 0epoch_student_loss = 0for batch_idx, (data, target) in enumerate(train_loader):data, target = data.to(device), target.to(device)# 获取教师模型的输出（logits）with torch.no_grad():teacher_logits = teacher_model(data)# 获取学生模型的输出student_logits = student_model(data)# 计算蒸馏损失total_loss, distill_loss, student_loss = distillation_criterion(student_logits, teacher_logits, target)# 反向传播和优化optimizer.zero_grad()total_loss.backward()optimizer.step()# 累计损失epoch_total_loss += total_loss.item()epoch_distill_loss += distill_loss.item()epoch_student_loss += student_loss.item()if batch_idx % 100 == 0:print(f'Student Training: Epoch {epoch+1}/{epochs} [{batch_idx*len(data)}/{len(train_loader.dataset)} 'f'({100.*batch_idx/len(train_loader):.0f}%)]\tTotal Loss: {total_loss.item():.6f} 'f'(Distill: {distill_loss.item():.6f}, Student: {student_loss.item():.6f})')# 记录平均训练损失history['total_loss'].append(epoch_total_loss / len(train_loader))history['distill_loss'].append(epoch_distill_loss / len(train_loader))history['student_loss'].append(epoch_student_loss / len(train_loader))# 在测试集上评估student_model.eval()test_loss = 0correct = 0with torch.no_grad():for data, target in test_loader:data, target = data.to(device), target.to(device)output = student_model(data)test_loss += F.cross_entropy(output, target, reduction='sum').item()pred = output.argmax(dim=1, keepdim=True)correct += pred.eq(target.view_as(pred)).sum().item()test_loss /= len(test_loader.dataset)accuracy = 100. * correct / len(test_loader.dataset)# 记录测试损失和准确率history['test_loss'].append(test_loss)history['test_accuracy'].append(accuracy)print(f'Student Test: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)')return student_model, history# 为了演示温度参数的效果，定义一个函数来显示不同温度下的概率分布
def visualize_temperature_effect(logits, temperatures=[1, 2, 5, 10]):plt.figure(figsize=(12, 6))for i, temp in enumerate(temperatures):probs = F.softmax(torch.tensor(logits) / temp, dim=0).numpy()plt.subplot(1, len(temperatures), i+1)plt.bar(range(len(probs)), probs)plt.title(f'T = {temp}')plt.ylim(0, 1)if i == 0:plt.ylabel('Probability')plt.xlabel('Class')plt.tight_layout()plt.savefig('temperature_effect.png')plt.close()# 比较不同温度参数下的知识蒸馏效果
def compare_temperatures(teacher_model, train_loader, test_loader, temperatures=[1, 2, 4, 8], epochs=3):device = torch.device("cuda" if torch.cuda.is_available() else "cpu")teacher_model.to(device)results = {}for temp in temperatures:print(f"\n--- Training student with temperature = {temp} ---")# 创建一个新的学生模型student_model = StudentModel()# 使用知识蒸馏训练学生模型student, history = train_student_with_distillation(teacher_model, student_model, train_loader, test_loader, temperature=temp, alpha=0.5, epochs=epochs)# 记录结果results[temp] = {'model': student,'history': history,'final_accuracy': history['test_accuracy'][-1]}# 比较不同温度下的测试准确率plt.figure(figsize=(10, 6))for temp in temperatures:plt.plot(results[temp]['history']['test_accuracy'], label=f'T = {temp}')plt.title('Test Accuracy with Different Temperatures')plt.xlabel('Epoch')plt.ylabel('Accuracy (%)')plt.legend()plt.grid(True)plt.savefig('temperature_comparison.png')plt.close()# 打印最终准确率print("\nFinal Accuracy with Different Temperatures:")for temp in temperatures:print(f"T = {temp}: {results[temp]['final_accuracy']:.2f}%")return results# 主函数
def main():print("Loading MNIST dataset...")train_loader, test_loader = load_data()# 演示温度参数对概率分布的影响print("\nVisualizing temperature effect...")example_logits = [10, 5, 2, 1, 0, -1, -2, -3, -4, -5]  # 10个类别的logitsvisualize_temperature_effect(example_logits)# 训练教师模型print("\nTraining teacher model...")teacher_model = TeacherModel()teacher_model = train_teacher(teacher_model, train_loader, test_loader, epochs=1)# 比较不同温度参数的效果print("\nComparing different temperature parameters...")results = compare_temperatures(teacher_model, train_loader, test_loader, temperatures=[1, 2, 4, 8], epochs=1)print("\nKnowledge Distillation with Temperature Mechanism completed!")if __name__ == '__main__':main()

5. 温度参数的选择与影响

温度参数的选择对知识蒸馏的效果有显著影响。一般来说：

• 过低的温度（接近1）：软标签接近于硬标签，无法充分传递教师模型中的"暗示性知识"
• 过高的温度（如10以上）：所有类别的概率趋于均匀，可能会丢失有用的信息
• 适中的温度（通常在2~5之间）：往往能达到较好的效果，既保留了类别间的相对关系，又能突显次要类别的信息

下面是不同温度对蒸馏效果的影响分析表：

温度	优点	缺点	适用场景
低温(T≈1)	保留强类别预测	难以传递细微差别	类别差异明显的简单任务
中温(T≈2-5)	平衡主次信息	需要更精细的调优	多数一般任务
高温(T>5)	强调模型的不确定性	可能过度平滑	类别相似性高的复杂任务

6. 知识蒸馏中温度调节的实验结果与分析

为了更直观地理解温度参数的影响，我们可以分析一些实验结果。以CIFAR-10数据集为例，使用ResNet-34作为教师模型，ResNet-18作为学生模型，不同温度下的蒸馏效果如下：

温度	学生模型准确率	模型大小减少	推理速度提升
无蒸馏	92.1%	0%	0%
T=1	92.3%	54%	32%
T=2	93.4%	54%	32%
T=4	93.8%	54%	32%
T=8	93.5%	54%	32%
T=16	92.8%	54%	32%

从上表可以看出，T=4时蒸馏效果最好，甚至超过了原始的教师模型。这是因为软标签提供了更丰富的信息，帮助学生模型更好地泛化。而当温度过高时（如T=16），效果反而下降，这可能是因为过度"软化"使得类别间的区分度降低了。

7. 温度调节的自适应方法

除了固定温度参数，一些研究还提出了自适应温度调节的方法：

1. 基于训练阶段的动态调整：在训练初期使用较高温度帮助学生模型快速学习分布，随着训练进行逐渐降低温度以聚焦于主要类别

2. 基于样本难度的自适应温度：对于简单样本使用较低温度，对于困难样本使用较高温度

3. 通过元学习优化温度参数：将温度作为可学习参数，通过验证集性能自动调整

以下是一个简单的自适应温度调节实现：

class AdaptiveDistillationLoss(nn.Module):def __init__(self, alpha=0.5, init_temperature=4.0, adapt_rate=0.01):super(AdaptiveDistillationLoss, self).__init__()self.alpha = alphaself.temperature = init_temperatureself.adapt_rate = adapt_ratedef forward(self, student_logits, teacher_logits, labels, epoch):# 根据训练进度调整温度current_temp = max(1.0, self.temperature * (1.0 - epoch * self.adapt_rate))# 使用当前温度计算软标签soft_targets = F.softmax(teacher_logits / current_temp, dim=1)soft_prob = F.log_softmax(student_logits / current_temp, dim=1)# 计算蒸馏损失distillation_loss = F.kl_div(soft_prob, soft_targets, reduction='batchmean') * (current_temp ** 2)# 计算学生模型与真实标签的交叉熵损失student_loss = F.cross_entropy(student_logits, labels)# 总损失total_loss = self.alpha * distillation_loss + (1 - self.alpha) * student_lossreturn total_loss, distillation_loss, student_loss, current_temp

8. 知识蒸馏的流程图

下面是知识蒸馏的完整流程图：

9. 温度参数对logits的影响可视化

下面是一个示例，展示了温度参数如何影响模型的logits：

import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn.functional as F# 示例logits
logits = torch.tensor([8.0, 4.0, 2.0, 0.0, -2.0])# 设置不同的温度值
temperatures = [0.5, 1.0, 2.0, 5.0, 10.0]plt.figure(figsize=(15, 10))# 对每个温度值计算softmax概率
for i, temp in enumerate(temperatures):# 应用带温度的softmaxprobs = F.softmax(logits / temp, dim=0).numpy()# 绘制概率分布plt.subplot(len(temperatures), 1, i+1)plt.bar(range(len(probs)), probs)plt.title(f'Temperature T = {temp}')plt.ylim(0, 1)plt.ylabel('Probability')plt.xticks(range(len(probs)), ['Class 1', 'Class 2', 'Class 3', 'Class 4', 'Class 5'])# 在柱子上标注概率值for j, prob in enumerate(probs):plt.text(j, prob + 0.02, f'{prob:.3f}', ha='center')plt.tight_layout()
plt.savefig('temperature_visualization.png')
plt.show()

10. 知识蒸馏温度参数的实际应用策略

在实际应用中，选择合适的温度参数通常需要结合具体任务和模型进行调试。以下是一些实用的策略：

1. 启发式温度选择：

   • 对于简单的分类任务，T=2~3通常就足够
   • 对于复杂的多类别任务，可以尝试T=4~6
   • 对于细粒度分类问题，更高的温度(T=6~10)可能有帮助

2. 基于模型结构的温度选择：

   • 模型容量差距大时（如ResNet-152到MobileNet），使用更高的温度
   • 模型架构相似时（如ResNet-50到ResNet-18），中等温度即可

3. 温度参数的网格搜索：

   • 可以通过在验证集上进行网格搜索（如T∈{1,2,4,8,16}）来找到最佳温度
   • 搜索时可以固定其他超参数，如α=0.5

4. 温度与α的联合调优：

   • 温度T和α参数（控制蒸馏损失与学生损失的权重）通常相互关联
   • 较高的温度可能需要较小的α值

结论

知识蒸馏中的温度调节机制是一个强大而精妙的技术，通过调整softmax函数的温度参数，我们可以控制概率分布的"软硬程度"，从而更有效地将教师模型的知识迁移到学生模型中。适当的温度设置可以帮助学生模型捕获到教师模型对不同类别的细微判断，而不仅仅是最终的硬标签。

在实践中，温度参数的选择需要根据具体任务和模型架构进行调整。通常，中等的温度值（T=2~5）在大多数情况下能够取得较好的效果，但对于特定任务，可能需要通过实验来寻找最优值。同时，自适应温度调节方法也提供了更灵活的选择，可以根据训练阶段或样本特性动态调整温度参数。

接下来，我们将在第二部分探讨模型剪枝技术，特别是通道剪枝与权重剪枝的精度损失差异。

---

清华大学全五版的《DeepSeek教程》完整的文档需要的朋友，关注我私信：deepseek 即可获得。

怎么样今天的内容还满意吗？再次感谢朋友们的观看，关注GZH：凡人的AI工具箱，回复666，送您价值199的AI大礼包。最后，祝您早日实现财务自由，还请给个赞，谢谢！