📚 课程目录第2章知识蒸馏基础实现

2. 知识蒸馏基础实现

任务目标

实现一个完整的知识蒸馏训练流程：使用预训练的ResNet-34作为教师模型，指导学生模型ResNet-18在CIFAR-10数据集上进行训练。通过软标签（soft label）与温度缩放将教师知识迁移至学生模型，最终使学生模型获得超越独立训练的精度。

环境依赖

Python 3.8+，PyTorch 1.10+，torchvision 0.11+
GPU（推荐）或CPU均可运行，GPU训练约15分钟，CPU约2小时
无需额外安装包，CIFAR-10由torchvision自动下载

核心实现

以下代码完整实现了知识蒸馏训练闭环：教师模型加载预训练权重并冻结，学生模型从零训练；蒸馏损失由KL散度（软标签）与交叉熵（硬标签）加权组合。训练完成后自动评估学生模型在测试集上的准确率。

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms

# 设备配置
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 数据加载（CIFAR-10）
transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])

trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2)
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)
testloader = torch.utils.data.DataLoader(testset, batch_size=256, shuffle=False, num_workers=2)

# 教师模型：ResNet-34（预训练）
teacher = torchvision.models.resnet34(weights=torchvision.models.ResNet34_Weights.IMAGENET1K_V1)
teacher.fc = nn.Linear(teacher.fc.in_features, 10)  # 适配CIFAR-10
teacher = teacher.to(device)
teacher.eval()
for param in teacher.parameters():
    param.requires_grad = False

# 学生模型：ResNet-18（随机初始化）
student = torchvision.models.resnet18(weights=None)
student.fc = nn.Linear(student.fc.in_features, 10)
student = student.to(device)

# 蒸馏超参数
T = 4.0          # 温度
alpha = 0.7      # 软标签损失权重
epochs = 30
lr = 0.05

criterion_ce = nn.CrossEntropyLoss()
criterion_kl = nn.KLDivLoss(reduction='batchmean')
optimizer = optim.SGD(student.parameters(), lr=lr, momentum=0.9, weight_decay=5e-4)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)

def train_one_epoch():
    student.train()
    total_loss = 0.0
    for images, labels in trainloader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()

        with torch.no_grad():
            teacher_logits = teacher(images)

        student_logits = student(images)

        # 软标签损失（温度缩放后KL散度）
        loss_soft = criterion_kl(
            F.log_softmax(student_logits / T, dim=1),
            F.softmax(teacher_logits / T, dim=1)
        ) * (T * T)

        # 硬标签损失
        loss_hard = criterion_ce(student_logits, labels)

        loss = alpha * loss_soft + (1 - alpha) * loss_hard
        loss.backward()
        optimizer.step()
        total_loss += loss.item()

    return total_loss / len(trainloader)

def evaluate(model):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in testloader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    return 100.0 * correct / total

# 训练循环
print("开始知识蒸馏训练...")
for epoch in range(epochs):
    loss = train_one_epoch()
    scheduler.step()
    if (epoch + 1) % 5 == 0:
        acc = evaluate(student)
        print(f"Epoch [{epoch+1}/{epochs}]  Loss: {loss:.4f}  学生模型测试准确率: {acc:.2f}%")

# 最终评估
final_acc = evaluate(student)
print(f"\n训练完成。学生模型最终测试准确率: {final_acc:.2f}%")

运行与验证

运行方式：将上述代码保存为distill.py，在终端执行python distill.py。首次运行会自动下载CIFAR-10数据集（约170MB）和ResNet-34预训练权重（约83MB）。
通过条件：训练30个epoch后，学生模型（ResNet-18）测试准确率应达到85%以上。作为参考，独立训练的ResNet-18在相同配置下约82%～84%，蒸馏后应明显超过该基线。
验证要点：观察训练日志中每5个epoch打印的准确率是否持续上升；最终准确率是否超过85%。若低于80%，可检查温度T（调低至2～3）或alpha（调高至0.8～0.9）后重训。

关键说明

教师模型使用ImageNet预训练权重后替换最后一层全连接，由于CIFAR-10图像尺寸为32×32，ResNet-34的输入层会自动下采样适配，无需额外修改。
蒸馏损失中(T * T)缩放因子用于平衡梯度量级，这是标准做法，不可省略。
教师模型全程处于eval()模式且梯度关闭，确保只前向传播、不更新参数。
若显存不足（教师+学生同时前向），可将batch_size从128降至64。