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%(纯软标签)
运行与验证
- 准备教师模型:需先训练一个 ResNet-18 在 CIFAR-10 上,保存为
teacher_resnet18_cifar10.pth。若已有其他教师模型,替换加载路径即可。 - 执行训练:直接运行上述脚本,程序会依次对
alpha = 0.0, 0.3, 0.5, 0.7, 1.0各训练 3 个 epoch 并输出测试准确率。 - 判定通过条件:观察输出结果,应看到
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)以获得最佳性能。教师模型需提前训练好,本案例不包含教师训练过程。