3. 软标签与硬标签混合训练

任务目标

在知识蒸馏中,仅使用教师模型的软标签训练学生模型,可能因教师预测中的噪声而丢失真实类别信息。本案例实现软标签(教师输出)与硬标签(真实one-hot)的混合训练,通过权重参数 alpha 控制两者比例,观察不同 alpha 下学生模型在测试集上的准确率变化。

环境与依赖

需要 Python 3.8+、PyTorch 2.0+、torchvision、numpy。数据集使用 CIFAR-10,教师模型为预训练的 ResNet-18(若网络不可用,代码会自动下载)。

核心实现

以下代码完成:定义学生模型(小型CNN)、加载预训练教师模型、实现混合损失函数(alpha * KL(学生logits, 教师logits) + (1-alpha) * CrossEntropy(学生logits, 真实标签)),并在不同 alpha 下训练并评估。

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

# ---------- 学生模型:小型CNN ----------
class StudentCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = 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.fc = nn.Linear(128 * 4 * 4, 10)

    def forward(self, x):
        x = self.conv(x)
        x = x.view(x.size(0), -1)
        return self.fc(x)

# ---------- 数据加载 ----------
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))
])
train_set = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
test_set = torchvision.datasets.CIFAR10(root='./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)

# ---------- 加载预训练教师 (ResNet-18) ----------
teacher = torchvision.models.resnet18(weights=torchvision.models.ResNet18_Weights.DEFAULT)
teacher.fc = nn.Linear(512, 10)  # CIFAR-10 类别数
teacher.load_state_dict(torch.load('teacher_resnet18_cifar10.pth'))  # 需预先训练并保存
teacher.eval()
for p in teacher.parameters():
    p.requires_grad = False

# ---------- 混合损失 ----------
def mixed_loss(student_logits, teacher_logits, labels, alpha, T=4.0):
    # 软标签损失:KL散度(温度缩放后)
    soft_loss = nn.KLDivLoss(reduction='batchmean')(
        F.log_softmax(student_logits / T, dim=1),
        F.softmax(teacher_logits / T, dim=1)
    ) * (T * T)
    # 硬标签损失:交叉熵
    hard_loss = F.cross_entropy(student_logits, labels)
    return alpha * soft_loss + (1 - alpha) * hard_loss

# ---------- 训练与评估 ----------
def train_and_evaluate(alpha, epochs=5):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    student = StudentCNN().to(device)
    teacher.to(device)
    optimizer = torch.optim.Adam(student.parameters(), lr=1e-3)
    teacher.eval()

    for epoch in range(epochs):
        student.train()
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            optimizer.zero_grad()
            with torch.no_grad():
                t_logits = teacher(images)
            s_logits = student(images)
            loss = mixed_loss(s_logits, t_logits, labels, alpha)
            loss.backward()
            optimizer.step()

    # 评估
    student.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in test_loader:
            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

# ---------- 主流程:扫描不同 alpha ----------
if __name__ == '__main__':
    alpha_list = [0.0, 0.3, 0.5, 0.7, 1.0]
    results = {}
    for a in alpha_list:
        acc = train_and_evaluate(alpha=a, epochs=3)  # 减少epoch便于快速演示
        results[a] = acc
        print(f'alpha={a:.1f} -> 测试准确率: {acc:.2f}%')
    # 预期输出示例(实际值因随机性略有波动):
    # alpha=0.0 -> 约 65-70%(纯硬标签)
    # alpha=0.5 -> 约 70-75%(混合)
    # alpha=1.0 -> 约 68-73%(纯软标签)

运行与验证

  1. 准备教师模型:需先训练一个 ResNet-18 在 CIFAR-10 上,保存为 teacher_resnet18_cifar10.pth。若已有其他教师模型,替换加载路径即可。
  2. 执行训练:直接运行上述脚本,程序会依次对 alpha = 0.0, 0.3, 0.5, 0.7, 1.0 各训练 3 个 epoch 并输出测试准确率。
  3. 判定通过条件:观察输出结果,应看到 alpha=0.5 附近的准确率不低于 alpha=0.0(纯硬标签)和 alpha=1.0(纯软标签),说明混合训练有效结合了两者优势。若 alpha=1.0 准确率明显低于 alpha=0.0,可适当增加训练轮次或调低温度 T 至 2~3 再验证。

关键说明

温度 T 控制软标签的平滑程度,T 越大类别分布越平滑。本实现中 alpha 为软标签损失权重,1-alpha 为硬标签权重。实际应用中建议在验证集上网格搜索 alpha(如 0.3~0.7 区间)和 T(2~6)以获得最佳性能。教师模型需提前训练好,本案例不包含教师训练过程。