4. 软标签与硬标签结合

任务说明

本案例实现知识蒸馏中软标签与硬标签损失的混合训练。教师模型为预训练ResNet18，学生模型为轻量CNN。通过加权求和软标签损失（KL散度）与硬标签损失（交叉熵），并调整权重系数α，观察学生模型在CIFAR-10验证集上的准确率变化。

环境依赖

需要 torch、torchvision、numpy。若本机无GPU，代码自动回退CPU运行（速度较慢，可将训练轮数调小验证流程）。

核心实现

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

# ---------- 模型定义 ----------
class TeacherNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.backbone = torchvision.models.resnet18(weights=None, num_classes=10)
    def forward(self, x):
        return self.backbone(x)

class StudentNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1), nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, 3, padding=1), nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(128*4*4, 256), nn.ReLU(),
            nn.Linear(256, 10)
        )
    def forward(self, x):
        return self.classifier(self.features(x))

# ---------- 数据加载 ----------
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))
])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = DataLoader(trainset, batch_size=128, shuffle=True)
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
testloader = DataLoader(testset, batch_size=256, shuffle=False)

# ---------- 训练准备 ----------
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
teacher = TeacherNet().to(device)
student = StudentNet().to(device)

# 加载预训练教师权重（此处用随机权重演示流程，实际使用请替换为训练好的权重）
# 正式实验请加载真实教师模型：teacher.load_state_dict(torch.load('teacher.pth'))
teacher.eval()
for p in teacher.parameters():
    p.requires_grad = False

optimizer = torch.optim.Adam(student.parameters(), lr=1e-3)
criterion_ce = nn.CrossEntropyLoss()
criterion_kl = nn.KLDivLoss(reduction='batchmean')
T = 4.0          # 温度
alpha = 0.7      # 软标签损失权重（硬标签权重为 1-alpha）

# ---------- 训练循环 ----------
def train_one_epoch():
    student.train()
    total_loss = 0.0
    for images, labels in trainloader:
        images, labels = images.to(device), labels.to(device)
        with torch.no_grad():
            teacher_logits = teacher(images)
        student_logits = student(images)

        # 硬标签损失
        loss_hard = criterion_ce(student_logits, labels)
        # 软标签损失（温度缩放后KL散度）
        loss_soft = criterion_kl(
            F.log_softmax(student_logits / T, dim=1),
            F.softmax(teacher_logits / T, dim=1)
        ) * (T * T)
        # 混合损失
        loss = alpha * loss_soft + (1 - alpha) * loss_hard

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    return total_loss / len(trainloader)

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

# ---------- 执行训练 ----------
if __name__ == '__main__':
    epochs = 5
    for epoch in range(epochs):
        loss = train_one_epoch()
        acc = evaluate()
        print(f'Epoch {epoch+1}/{epochs} | Loss: {loss:.4f} | Val Acc: {acc:.2f}%')
    print('训练完成')

运行与验证

直接执行脚本即可开始训练。代码会先下载CIFAR-10数据集（首次运行约需数分钟），随后进行5轮训练并输出每轮验证准确率。

验证通过条件：

• 程序无报错完整跑完5轮训练；
• 每轮输出的Loss值呈下降趋势；
• 最终验证准确率高于随机猜测（10%），且随训练轮数增加准确率逐步提升。

如需对比不同α值的影响，可修改代码中alpha = 0.7为其他值（如0.3、0.5、0.9）后重新运行，观察验证准确率变化。

关键说明

本案例使用随机初始化的教师模型仅用于演示蒸馏流程可运行性。实际应用中，教师模型需预先在数据集上训练至较高准确率，并将权重保存后在此处加载（替换代码中注释部分）。温度T控制软标签的平滑程度，α控制两类损失的相对重要性，二者均为可调超参数。