知识蒸馏实战：BCKD与CWD结合的新手入门指南

1次阅读

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

知识蒸馏（Knowledge Distillation）是一种模型压缩技术，通过让小型学生模型模仿大型教师模型的行为，实现性能提升。传统方法如 KD（Knowledge Distillation）仅使用教师模型的输出作为监督信号，而 BCKD 和 CWD 则提供了更精细的知识迁移方式。

BCKD（Bidirectional Collaborative Knowledge Distillation）：通过双向监督机制，让教师和学生模型相互学习，提升两者的表现。
CWD（Channel-wise Knowledge Distillation）：通过通道注意力机制，让学生模型学习教师模型的特征通道分布，实现更精细的特征对齐。

与传统蒸馏方法相比，BCKD 和 CWD 的结合在性能和复杂度上具有显著优势：

KD：仅使用教师模型的输出概率，计算简单但信息量有限。
FitNets：通过中间层特征对齐提升性能，但计算复杂度较高。
BCKD+CWD：结合双向监督和通道注意力，在保持较低复杂度的同时实现更高的精度。

BCKD 的核心思想是让学生和教师模型相互监督。具体实现包括以下步骤：

计算学生模型和教师模型的输出概率分布。
使用 KL 散度（Kullback-Leibler Divergence）衡量两者分布的差异。
将双向 KL 散度损失加权求和，作为总损失的一部分。

CWD 通过通道注意力机制对齐学生和教师模型的特征图：

对教师和学生模型的中间层特征图进行通道归一化。
计算通道注意力权重，重点关注信息量丰富的通道。
使用均方误差（MSE）对齐学生和教师模型的通道分布。

将 BCKD 和 CWD 的损失函数加权求和，作为最终的蒸馏损失：

BCKD 损失：双向 KL 散度损失。
CWD 损失：通道注意力对齐损失。
总损失：分类损失 + α * BCKD 损失 + β * CWD 损失。

以下是一个完整的 PyTorch 实现示例，基于 CIFAR-10 数据集：

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

# 定义 BCKD 损失函数
class BCKDLoss(nn.Module):
    def __init__(self, temperature=4):
        super(BCKDLoss, self).__init__()
        self.temperature = temperature
        self.kl_div = nn.KLDivLoss(reduction='batchmean')

    def forward(self, student_logits, teacher_logits):
        # 计算双向 KL 散度
        student_probs = torch.softmax(student_logits / self.temperature, dim=1)
        teacher_probs = torch.softmax(teacher_logits / self.temperature, dim=1)
        loss_student = self.kl_div(torch.log(student_probs), teacher_probs)
        loss_teacher = self.kl_div(torch.log(teacher_probs), student_probs)
        return (loss_student + loss_teacher) / 2

# 定义 CWD 损失函数
class CWDLoss(nn.Module):
    def __init__(self):
        super(CWDLoss, self).__init__()
        self.mse = nn.MSELoss()

    def forward(self, student_feats, teacher_feats):
        # 通道归一化
        student_norm = torch.norm(student_feats, p=2, dim=1, keepdim=True)
        teacher_norm = torch.norm(teacher_feats, p=2, dim=1, keepdim=True)
        student_normalized = student_feats / (student_norm + 1e-6)
        teacher_normalized = teacher_feats / (teacher_norm + 1e-6)
        return self.mse(student_normalized, teacher_normalized)

# 训练流程（简化版）def train(model, teacher_model, train_loader, optimizer, alpha=0.5, beta=0.5):
    model.train()
    teacher_model.eval()
    criterion_cls = nn.CrossEntropyLoss()
    criterion_bckd = BCKDLoss()
    criterion_cwd = CWDLoss()

    for data, target in train_loader:
        optimizer.zero_grad()
        output, feats = model(data)
        with torch.no_grad():
            teacher_output, teacher_feats = teacher_model(data)

        # 计算各项损失
        loss_cls = criterion_cls(output, target)
        loss_bckd = criterion_bckd(output, teacher_output)
        loss_cwd = criterion_cwd(feats, teacher_feats)
        total_loss = loss_cls + alpha * loss_bckd + beta * loss_cwd

        total_loss.backward()
        optimizer.step()

在 CIFAR-10 数据集上的实验结果显示：