CKD知识蒸馏实战指南：从模型压缩到部署优化的完整流程

1次阅读

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

最近在部署一个图像分类模型时，遇到了显存爆炸和推理延迟的问题。原本在实验室跑得好好的 ResNet50，一到生产环境就「水土不服」。这让我开始研究模型压缩技术，发现知识蒸馏（Knowledge Distillation）尤其是 CKD（Collaborative Knowledge Distillation）在精度和效率的平衡上表现突出。

先说说常见的几种压缩方案：

量化：把 FP32 转成 INT8，简单粗暴但精度损失明显
剪枝：去掉不重要的神经元，但需要复杂的重训练
蒸馏：让大模型（教师）教小模型（学生），保留「知识」而非简单模仿

而 CKD 的创新点在于： 双教师协同教学 。就像一个学生同时向语文老师和数学老师学习，能吸收不同视角的知识。

graph LR
    A[教师模型 A] -->| 软标签 + 特征 | C[学生模型]
    B[教师模型 B] -->| 软标签 + 特征 | C
    C -->|KL 散度 | D[损失函数]

与传统单教师蒸馏不同，CKD 的梯度来自两个方向：

教师 A 提供的类别间关系（通过高温 softmax 获得）
教师 B 提供的中间层特征分布

温度系数 $\tau$ 控制着软标签的「模糊程度」：

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

当 $\tau\to\infty$ 时，所有类别的概率趋近相同；当 $\tau\to0$ 时，退化成普通 one-hot 标签。实践中发现 $\tau=3\sim5$ 时效果最佳。

CKD 的损失包含三部分：

$$
\mathcal{L} = \alpha\mathcal{L}{KL} + \beta\mathcal{L}
$$} + \gamma\mathcal{L}_{task

其中跨层特征对齐的 L2 损失最值得关注：

def feature_loss(student_feats, teacher_feats):
    # 对每个特征图进行 L2 归一化
    s_norm = [F.normalize(f, dim=1) for f in student_feats]
    t_norm = [F.normalize(f, dim=1) for f in teacher_feats]

    # 逐层计算 MSE
    return sum([F.mse_loss(s, t) for s,t in zip(s_norm, t_norm)])

传统 KLDivLoss 对异常值敏感，改进版加入动态权重：

class DynamicKLoss(nn.Module):
    def __init__(self, tau=4.0, reduction='batchmean'):
        super().__init__()
        self.tau = tau
        self.reduction = reduction

    def forward(self, student_logits, teacher_logits):
        # 计算软标签
        s_probs = F.softmax(student_logits/self.tau, dim=1)
        t_probs = F.softmax(teacher_logits/self.tau, dim=1)

        # 根据置信度动态调整权重
        confidence = 1.0 - (t_probs.max(dim=1)[0] / self.tau)
        loss = F.kl_div(s_probs.log(), 
            t_probs.detach(),
            reduction='none'
        ).sum(dim=1) * confidence

        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'batchmean':
            return loss.sum() / student_logits.size(0)

MixUp 能显著提升蒸馏效果，但要注意对齐教师和学生的输入：

def mixup_data(x, y, alpha=0.4):
    if alpha > 0:
        lam = np.random.beta(alpha, alpha)
    else:
        lam = 1

    batch_size = x.size(0)
    index = torch.randperm(batch_size).to(x.device)

    mixed_x = lam * x + (1 - lam) * x[index]
    y_a, y_b = y, y[index]
    return mixed_x, y_a, y_b, lam

# 使用时需要同时处理三个损失
inputs, targets_a, targets_b, lam = mixup_data(x, y)
outputs = student(inputs)
loss = lam * criterion(outputs, targets_a) + \
       (1 - lam) * criterion(outputs, targets_b)

获取中间层特征的神器：

features = {}
def get_features(name):
    def hook(model, input, output):
        features[name] = output.detach()
    return hook

# 在教师模型上注册 hook
teacher_layer = teacher_model.layer3[-1].conv2
handle = teacher_layer.register_forward_hook(get_features('teacher_feats'))

# 前向传播后即可获取特征
_ = teacher_model(inputs)
teacher_feats = features['teacher_feats']