知识蒸馏技术解析：从bckd原理解析到模型压缩实战

1次阅读

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

随着深度学习模型规模的不断增大，模型压缩成为了工业界和学术界共同关注的焦点。知识蒸馏（Knowledge Distillation, KD）作为一种有效的模型压缩技术，能够将大型教师模型的知识迁移到小型学生模型中，从而在保持较高精度的同时大幅减少模型参数量和计算量。

知识蒸馏技术自 2015 年由 Hinton 等人提出以来，已经发展出多种变体，其中 bckd（Bi-directional Contrastive Knowledge Distillation）是一种较新的方法，它通过双向对比学习的方式，更好地捕捉教师模型和学生模型之间的知识差异。

传统知识蒸馏主要依赖于教师模型输出的软标签（soft targets）来指导学生模型的训练，其损失函数通常包括两部分：学生模型输出与真实标签的交叉熵损失，以及学生模型输出与教师模型输出的 KL 散度损失。数学表达式如下：

$$
L_{KD} = \alpha \cdot L_{CE}(y, \sigma(z_s)) + (1-\alpha) \cdot T^2 \cdot L_{KL}(\sigma(z_t/T), \sigma(z_s/T))
$$

其中，$z_t$ 和 $z_s$ 分别表示教师模型和学生模型的 logits，$T$ 是温度系数，$\alpha$ 是权重参数，$\sigma$ 表示 softmax 函数。

而 bckd 则引入了对比学习的思想，通过构建正负样本对，让学生模型不仅学习教师模型的输出分布，还能学习到特征空间中的相似性关系。其核心损失函数可以表示为：

$$
L_{bckd} = L_{KD} + \beta \cdot L_{contrastive}
$$

其中，$L_{contrastive}$ 是对比损失，通常采用 InfoNCE 损失函数形式：

$$
L_{contrastive} = -\log \frac{\exp(sim(f_t, f_s)/\tau)}{\sum_{i=1}^N \exp(sim(f_t, f_{s,i})/\tau)}
$$

这里，$f_t$ 和 $f_s$ 分别表示教师模型和学生模型的特征表示，$\tau$ 是温度参数，$N$ 是负样本数量。

在实现 bckd 时，首先需要选择合适的教师模型和学生模型架构。教师模型通常是一个预训练好的大型模型（如 ResNet50），而学生模型则是一个轻量级模型（如 MobileNetV2）。

import torch
import torch.nn as nn
import torchvision.models as models

# 教师模型
teacher_model = models.resnet50(pretrained=True)
teacher_model.eval()  # 固定教师模型参数

# 学生模型
student_model = models.mobilenet_v2(pretrained=False)

bckd 中有几个关键超参数需要仔细调优：

温度系数 T：控制软标签的平滑程度，通常取值在 1 -10 之间
权重参数 α 和 β：平衡不同损失项的重要性
对比损失温度 τ：控制对比学习的难度

import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

def train_bckd(teacher, student, train_loader, epochs, device):
    # 定义损失函数
    criterion_kd = nn.KLDivLoss(reduction='batchmean')
    criterion_ce = nn.CrossEntropyLoss()
    criterion_contrastive = nn.CrossEntropyLoss()  # 用于对比学习

    optimizer = optim.Adam(student.parameters(), lr=0.001)

    # 温度参数
    T = 3.0
    alpha = 0.5
    beta = 1.0

    for epoch in range(epochs):
        student.train()
        total_loss = 0.0

        for inputs, labels in train_loader:
            inputs, labels = inputs.to(device), labels.to(device)

            # 前向传播
            with torch.no_grad():
                teacher_logits, teacher_features = teacher(inputs, return_features=True)

            student_logits, student_features = student(inputs, return_features=True)

            # 计算各项损失
            # 1. 传统 KD 损失
            loss_kd = criterion_kd(F.log_softmax(student_logits/T, dim=1),
                F.softmax(teacher_logits/T, dim=1)
            ) * (T * T)

            # 2. 分类损失
            loss_ce = criterion_ce(student_logits, labels)

            # 3. 对比损失
            # 这里简化实现，实际应用中需要构建正负样本对
            loss_contrastive = criterion_contrastive(F.normalize(student_features, dim=1),
                F.normalize(teacher_features, dim=1)
            )

            # 总损失
            loss = alpha * loss_ce + (1-alpha) * loss_kd + beta * loss_contrastive

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

            total_loss += loss.item()

        print(f'Epoch {epoch+1}, Loss: {total_loss/len(train_loader):.4f}')

    return student

我们在 CIFAR-10 数据集上对比了传统 KD 和 bckd 的性能。实验设置如下：