📚 课程目录第2章温度参数调优实验

2. 温度参数调优实验

任务目标

本实验在MNIST数据集上训练一个教师模型（5层全连接网络），随后用不同温度值 T ∈ {1, 2, 4, 8, 16} 分别蒸馏训练学生模型（3层全连接网络），对比各温度下学生模型的测试准确率与软标签熵，确定最佳蒸馏温度。

环境与依赖

需要 Python 3.8+、PyTorch 2.0+、torchvision、numpy、matplotlib。MNIST数据集首次运行时会自动下载至 ./data 目录。

核心实现

以下代码完整实现教师训练、不同温度下的蒸馏训练、软标签熵计算与结果对比。蒸馏损失采用 KL散度(学生logits/T, 教师logits/T) * T² + CE(学生logits, 真实标签) 的标准形式。

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 numpy as np
import matplotlib.pyplot as plt

# ---------- 数据加载 ----------
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
train_set = datasets.MNIST('./data', train=True, download=True, transform=transform)
test_set = datasets.MNIST('./data', train=False, download=True, transform=transform)
train_loader = DataLoader(train_set, batch_size=128, shuffle=True)
test_loader = DataLoader(test_set, batch_size=256, shuffle=False)

# ---------- 网络定义 ----------
class TeacherNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Sequential(
            nn.Linear(784, 512), nn.ReLU(),
            nn.Linear(512, 256), nn.ReLU(),
            nn.Linear(256, 128), nn.ReLU(),
            nn.Linear(128, 64), nn.ReLU(),
            nn.Linear(64, 10)
        )
    def forward(self, x):
        return self.fc(x.view(x.size(0), -1))

class StudentNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Sequential(
            nn.Linear(784, 256), nn.ReLU(),
            nn.Linear(256, 128), nn.ReLU(),
            nn.Linear(128, 10)
        )
    def forward(self, x):
        return self.fc(x.view(x.size(0), -1))

# ---------- 训练与评估函数 ----------
def train_epoch(model, loader, optimizer, criterion):
    model.train()
    total_loss, correct = 0, 0
    for data, target in loader:
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * data.size(0)
        correct += (output.argmax(1) == target).sum().item()
    return total_loss / len(loader.dataset), correct / len(loader.dataset)

def evaluate(model, loader):
    model.eval()
    correct = 0
    with torch.no_grad():
        for data, target in loader:
            output = model(data)
            correct += (output.argmax(1) == target).sum().item()
    return correct / len(loader.dataset)

def soft_label_entropy(model, loader, T):
    """计算模型在温度T下软标签的平均熵（以自然对数底）"""
    model.eval()
    entropies = []
    with torch.no_grad():
        for data, _ in loader:
            logits = model(data) / T
            probs = F.softmax(logits, dim=1)
            ent = -(probs * torch.log(probs + 1e-10)).sum(dim=1)
            entropies.append(ent.mean().item())
    return np.mean(entropies)

# ---------- 1. 训练教师模型 ----------
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
teacher = TeacherNet().to(device)
opt_t = optim.Adam(teacher.parameters(), lr=1e-3)
ce_loss = nn.CrossEntropyLoss()

print("训练教师模型...")
for epoch in range(5):
    train_epoch(teacher, train_loader, opt_t, ce_loss)
teacher_acc = evaluate(teacher, test_loader)
print(f"教师模型测试准确率: {teacher_acc:.4f}")

# ---------- 2. 不同温度下蒸馏训练学生模型 ----------
def distill_train(student, teacher, T, epochs=5):
    opt_s = optim.Adam(student.parameters(), lr=1e-3)
    ce = nn.CrossEntropyLoss()
    student.train()
    teacher.eval()
    for _ in range(epochs):
        for data, target in train_loader:
            data, target = data.to(device), target.to(device)
            opt_s.zero_grad()
            with torch.no_grad():
                t_logits = teacher(data)
            s_logits = student(data)
            loss_soft = F.kl_div(
                F.log_softmax(s_logits / T, dim=1),
                F.softmax(t_logits / T, dim=1),
                reduction='batchmean'
            ) * (T * T)
            loss_hard = ce(s_logits, target)
            loss = loss_soft + loss_hard
            loss.backward()
            opt_s.step()

temperatures = [1, 2, 4, 8, 16]
results = []

for T in temperatures:
    student = StudentNet().to(device)
    distill_train(student, teacher, T)
    acc = evaluate(student, test_loader)
    entropy = soft_label_entropy(student, test_loader, T)
    results.append((T, acc, entropy))
    print(f"T={T:2d} | 学生测试准确率: {acc:.4f} | 软标签平均熵: {entropy:.4f}")

# ---------- 3. 结果可视化 ----------
Ts = [r[0] for r in results]
accs = [r[1] for r in results]
ents = [r[2] for r in results]

fig, ax1 = plt.subplots(figsize=(8, 5))
ax1.plot(Ts, accs, 'o-', color='#d62728', label='Student Accuracy')
ax1.set_xlabel('Temperature T')
ax1.set_ylabel('Test Accuracy', color='#d62728')
ax1.tick_params(axis='y', labelcolor='#d62728')
ax1.set_xticks(Ts)

ax2 = ax1.twinx()
ax2.plot(Ts, ents, 's--', color='#1f77b4', label='Soft Label Entropy')
ax2.set_ylabel('Avg Soft Label Entropy', color='#1f77b4')
ax2.tick_params(axis='y', labelcolor='#1f77b4')

plt.title('Effect of Temperature on Distillation')
fig.tight_layout()
plt.savefig('temperature_analysis.png', dpi=120)
plt.show()

# 输出最佳温度
best_T = max(results, key=lambda x: x[1])
print(f"\n最佳蒸馏温度: T={best_T[0]}, 对应学生准确率: {best_T[1]:.4f}")

运行方式

将上述代码保存为 temperature_distill.py，在终端执行：

python temperature_distill.py

程序会自动下载MNIST数据集并依次完成教师训练、5组温度蒸馏实验，最终打印各温度下的准确率与熵值，并生成 temperature_analysis.png 曲线图。

验证方法与判定标准

运行成功标志：终端输出5行 T=... | 学生测试准确率: ... | 软标签平均熵: ...，且无异常报错。
实验有效性：教师模型测试准确率应高于任一学生模型（通常教师可达98%以上，学生约97%左右），否则说明教师未充分训练或蒸馏失效。
温度影响规律：软标签平均熵随T增大而单调上升（T=1时熵最小，T=16时熵最大），这验证了高温产生更平滑的概率分布。
最佳温度判定：比较5组学生准确率，最高值对应的T即为最佳温度。通常最佳T落在2~8区间（预期准确率约97.5%~98.2%），若T=1时准确率显著低于其他温度，说明软标签信息未被有效利用。

关键说明

蒸馏损失中 T² 系数用于平衡梯度尺度，使不同温度下的损失量级可比，这是标准做法。
软标签熵在测试集上计算，反映学生模型输出分布的平滑程度；熵越高说明类别间概率越接近，即软标签携带更多类间相似性信息。
若需复现精确结果，可设置随机种子 torch.manual_seed(0)，但不同温度间的相对趋势不会改变。