知识蒸馏实战：从BCKD公式解析到模型轻量化落地

1次阅读

没有评论

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

在现实场景中，大型深度学习模型（如 ResNet-50、BERT 等）虽然性能强大，但面临严重的部署难题：

算力需求高：大模型推理需要强大的 GPU 算力支持
内存占用大：模型参数可能达到数百 MB 甚至 GB 级别
延迟问题：在移动端或边缘设备上响应速度慢

传统知识蒸馏 (Knowledge Distillation, KD) 通过教师 - 学生模型的方式缓解这一问题，但存在明显局限：

仅使用软标签 (soft targets) 进行监督
忽略了中间层特征的结构化信息
特征对齐效果不理想，导致小模型精度损失严重

BCKD(Bilateral Contrastive Knowledge Distillation)通过引入对比学习机制，显著提升了特征迁移效率。其损失函数包含四个关键组件：

教师模型特征对比：
$$L_{tea} = -\frac{1}{N}\sum_{i=1}^N \log\frac{\exp(sim(z_i^t,z_j^t)/\tau)}{\sum_{k=1}^K \exp(sim(z_i^t,z_k^t)/\tau)}$$
学生模型特征对比：
$$L_{stu} = -\frac{1}{N}\sum_{i=1}^N \log\frac{\exp(sim(z_i^s,z_j^s)/\tau)}{\sum_{k=1}^K \exp(sim(z_i^s,z_k^s)/\tau)}$$
师生特征对齐：
$$L_{align} = -\frac{1}{N}\sum_{i=1}^N \log\frac{\exp(sim(z_i^t,z_i^s)/\tau)}{\sum_{k=1}^K \exp(sim(z_i^t,z_k^s)/\tau)}$$
最终目标函数：
$$L_{total} = \alpha L_{tea} + \beta L_{stu} + \gamma L_{align}$$

其中：
– $z^t,z^s$ 分别表示教师和学生特征
– $sim(\cdot)$ 为余弦相似度
– $\tau$ 是温度系数
– $\alpha,\beta,\gamma$ 为平衡权重

class ProjectionHead(nn.Module):
    """
    将特征映射到对比学习空间
    参数说明：in_dim: 输入特征维度（推荐：教师 2048/ 学生 512）proj_dim: 投影维度（推荐：256-1024）"""
    def __init__(self, in_dim=512, proj_dim=256):
        super().__init__()
        self.fc1 = nn.Linear(in_dim, proj_dim)
        self.ln1 = nn.LayerNorm(proj_dim)  # 稳定训练
        self.fc2 = nn.Linear(proj_dim, proj_dim)

    def forward(self, x):
        x = F.relu(self.ln1(self.fc1(x)))
        x = self.fc2(x)
        return F.normalize(x, p=2, dim=1)  # L2 归一化

def bckd_loss(tea_feat, stu_feat, temp=0.1, queue_size=65536):
    """
    计算 BCKD 三部分损失
    参数说明：temp: 温度系数（推荐 0.05-0.5）queue_size: 负样本队列大小（推荐 4096-65536）"""
    # 相似度矩阵计算
    sim_tt = torch.mm(tea_feat, tea_feat.T) / temp  # 教师 - 教师
    sim_ss = torch.mm(stu_feat, stu_feat.T) / temp  # 学生 - 学生
    sim_ts = torch.mm(tea_feat, stu_feat.T) / temp  # 师生对齐

    # 对比损失计算
    loss_tea = F.cross_entropy(sim_tt, labels)  # 教师特征对比
    loss_stu = F.cross_entropy(sim_ss, labels)  # 学生特征对比
    loss_align = F.cross_entropy(sim_ts, labels)  # 特征对齐

    return 0.3*loss_tea + 0.3*loss_stu + 0.4*loss_align  # 加权求和

# 在训练循环中加入：scaler = GradScaler()  # 混合精度训练
optimizer.zero_grad()

with autocast():  # 自动混合精度
    loss = bckd_loss(tea_feat, stu_feat)

scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)  # 梯度裁剪
scaler.step(optimizer)
scaler.update()