PyTorch实战MNIST知识蒸馏：从模型压缩到部署优化的全流程解析

1次阅读

共计 2688 个字符，预计需要花费 7 分钟才能阅读完成。

在边缘设备上部署深度学习模型时，我们常常面临两个主要问题：

计算资源限制 ：边缘设备通常内存有限，计算能力较弱，难以运行大型模型
能耗问题 ：复杂的模型会消耗更多电力，影响设备续航

传统 MNIST 分类模型虽然相对简单，但在资源极度受限的环境下（如 MCU），依然可能成为瓶颈。以一个典型的 5 层 CNN 为例，其参数规模约为 1.2MB，在 Cortex-M4 处理器上单次推断需要约 50ms – 这对于实时性要求高的应用仍然不够理想。

技术	压缩率	精度损失	硬件要求	实现复杂度
知识蒸馏	3-5x	<2%	低	中
剪枝	2-10x	1-5%	低	高
量化	4x	<1%	需支持 INT8	低

知识蒸馏在保持较高精度的同时，提供了不错的压缩率，且对硬件没有特殊要求，是边缘部署的理想选择。

# 教师模型 (ResNet18)
class TeacherModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.resnet = models.resnet18(num_classes=10)
        self.resnet.conv1 = nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1, bias=False)  # 适配 MNIST 单通道输入

    def forward(self, x):  # x: [B, 1, 28, 28]
        return self.resnet(x)

# 学生模型 (简化 CNN)
class StudentModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(nn.Conv2d(1, 16, 3, padding=1),  # [B,16,28,28]
            nn.ReLU(),
            nn.MaxPool2d(2),                # [B,16,14,14]
            nn.Conv2d(16, 32, 3, padding=1), # [B,32,14,14]
            nn.ReLU(),
            nn.MaxPool2d(2)                 # [B,32,7,7]
        )
        self.classifier = nn.Linear(32*7*7, 10)

    def forward(self, x):  # x: [B, 1, 28, 28]
        x = self.features(x)
        x = x.view(x.size(0), -1)
        return self.classifier(x)

知识蒸馏的核心是使用 ” 软标签 ”（soft targets），通过温度参数 τ 控制标签的 ” 软化 ” 程度：

$$
q_i = \frac{\exp(z_i/\tau)}{\sum_j \exp(z_j/\tau)}
$$

τ 值越大，分布越平滑；τ= 1 时退化为标准 softmax。实践中发现 τ =3- 5 对 MNIST 效果最佳。

def soft_target_loss(student_logits, teacher_logits, temp):
    """计算软化后的 KL 散度损失"""
    soft_teacher = F.softmax(teacher_logits / temp, dim=1)
    soft_student = F.log_softmax(student_logits / temp, dim=1)
    return F.kl_div(soft_student, soft_teacher, reduction='batchmean') * (temp**2)

总损失是硬标签交叉熵和软标签 KL 散度的加权和：

$$
L = \alpha L_{CE} + (1-\alpha)L_{KL}
$$

def distillation_loss(student_logits, teacher_logits, labels, temp, alpha):
    """
    参数:
        student_logits: 学生网络输出 [B,10]
        teacher_logits: 教师网络输出 [B,10]
        labels: 真实标签 [B]
        temp: 温度参数
        alpha: 硬标签损失权重
    """
    ce_loss = F.cross_entropy(student_logits, labels)
    kl_loss = soft_target_loss(student_logits, teacher_logits, temp)
    return alpha * ce_loss + (1-alpha) * kl_loss

当温度 τ 设置过小时，softmax 梯度可能变得非常陡峭，导致训练不稳定。解决方案：

初始使用较大 τ 值（如 10），随着训练逐步降低
添加梯度裁剪：torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

通过实验发现，学生网络的宽度与蒸馏效果存在非线性关系：

学生模型宽度	参数量	测试准确率
8-16-32	23K	97.2%
16-32-64	89K	98.1%
32-64-128	350K	98.3%

建议：先从教师模型 1 / 4 参数量开始尝试，逐步调整。

torch.onnx.export(
    student_model,
    dummy_input, 
    "student.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}},
    opset_version=11
)

常见问题：