4. 知识蒸馏应用

任务目标

在嵌入式设备上，大模型推理延迟高、内存占用大。本案例使用知识蒸馏将教师模型（大模型）的知识迁移到学生模型（小模型），并对比两者在相同嵌入式推理环境下的推理延迟与精度，验证蒸馏后的小模型能否在显著降低延迟的同时保持接近的精度。

环境与依赖

  • Python 3.8+，PyTorch 1.10+，torchvision，ONNX Runtime（用于延迟对比）
  • 嵌入式推理环境：本案例以ONNX Runtime CPU执行器模拟嵌入式推理（无GPU），实际部署时可替换为RKNN/TensorRT等目标平台运行时

实现方案

采用CIFAR-10分类任务：教师模型为ResNet-18，学生模型为轻量级CNN（参数量约为教师的1/20）。蒸馏损失 = 硬标签交叉熵 + α × KL散度(学生logits, 教师logits / T)。训练完成后导出ONNX，用ONNX Runtime测量单张图片平均推理延迟。

关键完整代码

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import onnxruntime as ort
import numpy as np
import time

# ---------- 模型定义 ----------
class TeacherNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = torchvision.models.resnet18(num_classes=10)
    def forward(self, x):
        return self.model(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.AdaptiveAvgPool2d(1)
        )
        self.classifier = nn.Linear(128, 10)
    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        return self.classifier(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 = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=256, shuffle=False)

# ---------- 训练教师模型（已预训练则跳过） ----------
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
teacher = TeacherNet().to(device)
teacher.load_state_dict(torch.load('teacher.pth'))  # 若不存在则先训练
# 训练代码（简化）：交叉熵训练30个epoch，保存为teacher.pth
# 此处假设teacher.pth已存在，实际运行时需先执行训练

# ---------- 蒸馏训练学生模型 ----------
student = StudentNet().to(device)
teacher.eval()
criterion_ce = nn.CrossEntropyLoss()
criterion_kl = nn.KLDivLoss(reduction='batchmean')
optimizer = optim.Adam(student.parameters(), lr=1e-3)
T = 4.0
alpha = 0.7

for epoch in range(20):
    student.train()
    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_ce = criterion_ce(student_logits, labels)
        loss_kl = criterion_kl(
            nn.functional.log_softmax(student_logits / T, dim=1),
            nn.functional.softmax(teacher_logits / T, dim=1)
        ) * (T * T)
        loss = alpha * loss_kl + (1 - alpha) * loss_ce
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

# ---------- 精度评估 ----------
def evaluate(model, loader):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in loader:
            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

teacher_acc = evaluate(teacher, testloader)
student_acc = evaluate(student, testloader)
print(f'教师精度: {teacher_acc:.2f}% | 学生精度: {student_acc:.2f}%')

# ---------- ONNX导出与延迟对比 ----------
def export_onnx(model, path):
    model.eval()
    dummy = torch.randn(1, 3, 32, 32).to(device)
    torch.onnx.export(model, dummy, path, input_names=['input'], output_names=['output'], opset_version=11)

export_onnx(teacher, 'teacher.onnx')
export_onnx(student, 'student.onnx')

def measure_latency(onnx_path, runs=100):
    sess = ort.InferenceSession(onnx_path, providers=['CPUExecutionProvider'])
    input_data = np.random.randn(1, 3, 32, 32).astype(np.float32)
    # 预热
    for _ in range(10):
        sess.run(None, {'input': input_data})
    times = []
    for _ in range(runs):
        start = time.perf_counter()
        sess.run(None, {'input': input_data})
        times.append((time.perf_counter() - start) * 1000)
    return np.mean(times)

teacher_latency = measure_latency('teacher.onnx')
student_latency = measure_latency('student.onnx')
print(f'教师平均延迟: {teacher_latency:.2f} ms | 学生平均延迟: {student_latency:.2f} ms')
print(f'延迟降低: {(1 - student_latency/teacher_latency)*100:.1f}%')

运行方式

  1. 首次运行需先训练教师模型：在代码中取消注释训练部分，训练30个epoch后保存为teacher.pth，再执行蒸馏。
  2. 直接运行上述完整脚本，会自动下载CIFAR-10数据集、加载教师模型、蒸馏训练学生模型、导出ONNX并测量延迟。

验证方法与通过条件

  • 精度对比：学生模型精度应达到教师精度的90%以上（例如教师92%，学生≥83%），且明显高于从头训练同结构小模型（通常低3-5个百分点）。
  • 延迟对比：学生模型ONNX在CPU上的平均推理延迟应显著低于教师模型（通常降低50%以上），且两者在同一运行环境下测量。
  • 判定标准：控制台输出教师/学生精度与延迟数据，满足上述两个条件即视为蒸馏成功。

常见问题

  • 若教师模型未预训练，需先完成训练步骤，否则加载teacher.pth会报错。
  • 蒸馏温度T建议范围2~6，α建议0.5~0.8，可根据精度表现调整。
  • 实际嵌入式部署时，将ONNX转换为目标平台格式（如RKNN），延迟数据会随硬件变化，但相对趋势一致。
← 上一章 📖 返回目录 下一章 →