知识蒸馏入门指南：从bckd原理到模型轻量化实战

1次阅读

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

在深度学习领域，大模型（如 ResNet、BERT）虽然精度高，但参数量和计算开销巨大，难以部署到移动设备或边缘计算场景。知识蒸馏（Knowledge Distillation）技术的核心思想是让轻量级的小模型（学生模型）通过模仿大模型（教师模型）的行为来提升自身性能。

传统蒸馏方法（如 Hinton 提出的 KL 散度蒸馏）主要利用教师模型的输出概率分布作为软标签（soft targets）。而 bckd（Background Contrastive Knowledge Distillation）通过引入注意力转移机制，让学生模型不仅学习教师模型的输出，还能捕捉其内部特征的空间关系。

bckd 的核心创新点在于对比学习思想的引入。其损失函数由三部分组成：

传统蒸馏损失 ：使用 KL 散度衡量教师和学生模型输出的差异
$$L_{kd} = \tau^2 \cdot KL(\sigma(z_s/\tau) || \sigma(z_t/\tau))$$
背景对比损失 ：让同类样本的特征在嵌入空间中更接近
$$L_{con} = -\log\frac{\exp(sim(f_s,f_t)/\tau)}{\sum_{k=1}^K \exp(sim(f_s,f_k)/\tau)}$$
注意力转移损失 ：通过特征图的空间相关性传递知识
$$L_{at} = |A_s – A_t|_F^2$$

（示意图：教师模型与学生模型通过多层注意力特征交互）

以下是 bckd 的完整实现，包含关键组件注释：

import torch
import torch.nn as nn
import torch.nn.functional as F

class BCKDLoss(nn.Module):
    def __init__(self, temp=4.0, alpha=0.5):
        super().__init__()
        self.temp = temp
        self.alpha = alpha  # 平衡系数

    def forward(self, student_out, teacher_out, 
                student_feat, teacher_feat):
        # 传统蒸馏损失
        soft_loss = F.kl_div(F.log_softmax(student_out/self.temp, dim=1),
            F.softmax(teacher_out/self.temp, dim=1),
            reduction='batchmean') * (self.temp**2)

        # 特征对比损失
        student_feat = F.normalize(student_feat, p=2, dim=1)
        teacher_feat = F.normalize(teacher_feat, p=2, dim=1)
        sim_matrix = student_feat @ teacher_feat.T / self.temp
        con_loss = -torch.log(torch.diag(F.softmax(sim_matrix, dim=1))).mean()

        # 组合损失
        return (1-self.alpha)*soft_loss + self.alpha*con_loss

准备教师模型（需提前训练好）和学生模型

teacher = resnet50(pretrained=True).eval()
student = mobilenet_v2(pretrained=False)

# 冻结教师模型参数
for param in teacher.parameters():
    param.requires_grad = False

修改训练循环

optimizer = torch.optim.Adam(student.parameters(), lr=1e-3)
criterion = BCKDLoss(temp=4.0, alpha=0.7)

for images, labels in dataloader:
    # 前向传播
    with torch.no_grad():
        t_out, t_feat = teacher(images)
    s_out, s_feat = student(images)

    # 计算损失
    loss = criterion(s_out, t_out, s_feat, t_feat)

    # 反向传播
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

在 CIFAR-10 上的对比实验结果：