1. 教师-学生网络基础蒸馏

任务目标

本案例实现知识蒸馏中最基础的教师-学生范式：先训练一个参数量较大的教师模型，再训练一个小型学生模型，通过教师输出的软标签（带温度系数的概率分布）传递知识，使学生模型在CIFAR-10分类任务上达到接近教师模型的准确率。

环境与依赖

需要 Python 3.8+、PyTorch 2.0+、torchvision。数据集由 torchvision 自动下载，无需额外准备。

核心实现

以下代码完整实现教师与学生模型定义、训练流程、蒸馏损失计算及验证逻辑。教师模型使用较宽的ResNet-18，学生模型使用通道数减半的ResNet-18，两者结构一致但参数量不同。

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms, models

# ---------- 模型定义 ----------
def get_model(width_mult, num_classes=10):
    """返回通道数按width_mult缩放的ResNet-18"""
    model = models.resnet18(num_classes=num_classes)
    # 缩放第一个卷积层和所有BasicBlock的通道
    model.conv1 = nn.Conv2d(3, int(64*width_mult), 3, 1, 1, bias=False)
    model.bn1 = nn.BatchNorm2d(int(64*width_mult))
    model.layer1 = _scale_layer(model.layer1, width_mult)
    model.layer2 = _scale_layer(model.layer2, width_mult)
    model.layer3 = _scale_layer(model.layer3, width_mult)
    model.layer4 = _scale_layer(model.layer4, width_mult)
    model.fc = nn.Linear(int(512*width_mult), num_classes)
    return model

def _scale_layer(layer, width_mult):
    """缩放ResNet BasicBlock的卷积通道数"""
    for name, module in layer.named_children():
        if isinstance(module, models.resnet.BasicBlock):
            module.conv1 = nn.Conv2d(int(module.conv1.in_channels*width_mult),
                                     int(module.conv1.out_channels*width_mult),
                                     module.conv1.kernel_size, module.conv1.stride,
                                     module.conv1.padding, bias=False)
            module.conv2 = nn.Conv2d(int(module.conv2.in_channels*width_mult),
                                     int(module.conv2.out_channels*width_mult),
                                     module.conv2.kernel_size, module.conv2.stride,
                                     module.conv2.padding, bias=False)
            module.bn1 = nn.BatchNorm2d(int(module.bn1.num_features*width_mult))
            module.bn2 = nn.BatchNorm2d(int(module.bn2.num_features*width_mult))
            if module.downsample is not None:
                module.downsample = nn.Sequential(
                    nn.Conv2d(int(module.downsample[0].in_channels*width_mult),
                              int(module.downsample[0].out_channels*width_mult),
                              module.downsample[0].kernel_size, module.downsample[0].stride,
                              module.downsample[0].padding, bias=False),
                    nn.BatchNorm2d(int(module.downsample[1].num_features*width_mult))
                )
    return layer

# ---------- 蒸馏损失 ----------
def distillation_loss(student_logits, teacher_logits, labels, T=4.0, alpha=0.7):
    """软标签蒸馏损失 + 硬标签交叉熵"""
    soft_targets = F.softmax(teacher_logits / T, dim=1)
    student_soft = F.log_softmax(student_logits / T, dim=1)
    kd_loss = F.kl_div(student_soft, soft_targets, reduction='batchmean') * (T * T)
    ce_loss = F.cross_entropy(student_logits, labels)
    return alpha * kd_loss + (1 - alpha) * ce_loss

# ---------- 训练与验证 ----------
def train_one_epoch(model, loader, optimizer, device, teacher=None, T=4.0, alpha=0.7):
    model.train()
    total_loss, correct, total = 0, 0, 0
    for inputs, labels in loader:
        inputs, labels = inputs.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        if teacher is not None:
            with torch.no_grad():
                teacher_outputs = teacher(inputs)
            loss = distillation_loss(outputs, teacher_outputs, labels, T, alpha)
        else:
            loss = F.cross_entropy(outputs, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * inputs.size(0)
        correct += (outputs.argmax(1) == labels).sum().item()
        total += labels.size(0)
    return total_loss / total, correct / total

@torch.no_grad()
def evaluate(model, loader, device):
    model.eval()
    correct, total = 0, 0
    for inputs, labels in loader:
        inputs, labels = inputs.to(device), labels.to(device)
        outputs = model(inputs)
        correct += (outputs.argmax(1) == labels).sum().item()
        total += labels.size(0)
    return correct / total

# ---------- 主流程 ----------
def main():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    batch_size = 128
    epochs = 10

    # 数据加载
    transform = 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))
    ])
    train_set = datasets.CIFAR10('./data', train=True, download=True, transform=transform)
    test_set = datasets.CIFAR10('./data', train=False, download=True, transform=transform)
    train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=2)
    test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=2)

    # 教师模型（宽模型）
    teacher = get_model(width_mult=1.0).to(device)
    teacher_opt = optim.SGD(teacher.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
    teacher_sched = optim.lr_scheduler.CosineAnnealingLR(teacher_opt, T_max=epochs)

    print("训练教师模型...")
    for epoch in range(epochs):
        train_loss, train_acc = train_one_epoch(teacher, train_loader, teacher_opt, device)
        teacher_sched.step()
        test_acc = evaluate(teacher, test_loader, device)
        print(f"Epoch {epoch+1}/{epochs} | Train Acc: {train_acc:.4f} | Test Acc: {test_acc:.4f}")

    # 学生模型（窄模型）
    student = get_model(width_mult=0.5).to(device)
    student_opt = optim.SGD(student.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
    student_sched = optim.lr_scheduler.CosineAnnealingLR(student_opt, T_max=epochs)

    print("\n训练学生模型（使用教师软标签蒸馏）...")
    for epoch in range(epochs):
        train_loss, train_acc = train_one_epoch(student, train_loader, student_opt, device,
                                                teacher=teacher, T=4.0, alpha=0.7)
        student_sched.step()
        test_acc = evaluate(student, test_loader, device)
        print(f"Epoch {epoch+1}/{epochs} | Train Acc: {train_acc:.4f} | Test Acc: {test_acc:.4f}")

    # 对比：无蒸馏训练学生模型
    student_no_kd = get_model(width_mult=0.5).to(device)
    student_no_kd_opt = optim.SGD(student_no_kd.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
    student_no_kd_sched = optim.lr_scheduler.CosineAnnealingLR(student_no_kd_opt, T_max=epochs)

    print("\n训练学生模型（无蒸馏，仅硬标签）...")
    for epoch in range(epochs):
        train_loss, train_acc = train_one_epoch(student_no_kd, train_loader, student_no_kd_opt, device)
        student_no_kd_sched.step()
        test_acc = evaluate(student_no_kd, test_loader, device)
        print(f"Epoch {epoch+1}/{epochs} | Train Acc: {train_acc:.4f} | Test Acc: {test_acc:.4f}")

    # 最终对比
    teacher_acc = evaluate(teacher, test_loader, device)
    student_kd_acc = evaluate(student, test_loader, device)
    student_no_kd_acc = evaluate(student_no_kd, test_loader, device)
    print(f"\n最终测试准确率:")
    print(f"教师模型: {teacher_acc:.4f}")
    print(f"学生模型(蒸馏): {student_kd_acc:.4f}")
    print(f"学生模型(无蒸馏): {student_no_kd_acc:.4f}")

if __name__ == '__main__':
    main()

运行方式

将上述代码保存为 kd_basic.py，在项目目录下执行：

python kd_basic.py

首次运行会自动下载CIFAR-10数据集到 ./data 目录。使用GPU时训练约需5-10分钟（取决于显卡），CPU约需30-60分钟。

验证方法与判定标准

程序运行结束后会打印三个模型的最终测试准确率。判定蒸馏有效的标准为：

• 教师模型准确率显著高于无蒸馏学生模型（通常高5%以上）；
• 蒸馏学生模型准确率高于无蒸馏学生模型，且接近教师模型（差距在3%以内为优秀，5%以内为有效）。

若满足上述条件，说明软标签成功传递了教师模型的类别间相似性知识，学生模型性能得到提升。

关键参数说明

• 温度T=4.0：控制软标签的平滑程度。T越大，概率分布越平滑，类别间相似性信息越丰富；T=1时退化为硬标签。
• 权重α=0.7：蒸馏损失的权重。α越大，模型越关注教师软标签；α=0时完全忽略教师。
• 宽度倍率0.5：学生模型通道数为教师的一半，参数量约为教师的1/4，计算量显著降低。

常见问题

• 学生模型不收敛：尝试降低学习率至0.01，或增加训练轮数至20。
• 蒸馏效果不明显：增大温度T至6-8，或调整α至0.5-0.9区间搜索最优值。
• 显存不足：减小batch_size至64或32，教师模型推理时已使用no_grad，不会额外占用梯度内存。