知识蒸馏实战指南：如何用AT方法压缩深度学习模型

1次阅读

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

在深度学习领域，模型越大通常意味着性能越好，但随之而来的问题是模型体积庞大、推理速度慢，尤其对于移动端和边缘计算场景来说，这成为了一个严重的瓶颈。想象一下，一个几百 MB 甚至几个 GB 的模型，在手机上运行可能需要几秒钟甚至更长时间，这显然无法满足实时性要求。

目前主流的模型压缩技术主要有三种：

模型剪枝 ：通过去除网络中的冗余连接或神经元来减小模型大小
量化：将模型参数从浮点数转换为低精度表示（如 8 位整数）
知识蒸馏 ：用一个大型教师模型来指导一个小型学生模型的训练

在这三种方法中，知识蒸馏（Knowledge Distillation）因其能够保持较好的模型性能而受到广泛关注。而基于注意力转移（Attention Transfer, AT）的知识蒸馏方法，则进一步提升了知识迁移的效率。

传统的知识蒸馏主要使用教师模型的输出概率分布（软目标）来指导学生模型的训练。而 AT 蒸馏则更进一步，它利用了教师模型中间层的注意力图（Attention Maps）来指导学生模型。

标准知识蒸馏 ：
只利用教师模型的输出层信息
通过 KL 散度最小化教师和学生输出的差异
公式：L = α·L_CE + (1-α)·L_KD
AT 知识蒸馏 ：
额外利用教师模型的中间层注意力信息
通过最小化教师和学生注意力图的差异
公式：L = L_CE + β·L_AT

注意力机制（Attention Mechanism）是深度学习中的重要概念，它能够帮助模型聚焦于输入数据的关键部分。在 AT 蒸馏中，我们利用教师模型的注意力图来指导学生模型的学习。

具体来说，对于卷积神经网络，我们可以将特征图的激活值视为注意力图。通过计算这些注意力图之间的差异，我们可以定义一个额外的损失函数：

L_AT = ||A_T - A_S||^2

其中 A_T 和 A_S 分别表示教师和学生模型的注意力图。在实际实现中，我们通常会对特征图进行归一化处理。

下面我们用 PyTorch 实现一个完整的 AT 蒸馏训练流程。假设我们使用 ResNet 作为基础架构。

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

# 定义注意力转移损失
class ATLoss(nn.Module):
    def __init__(self, p=2):
        super(ATLoss, self).__init__()
        self.p = p

    def forward(self, fm_s, fm_t):
        # 对特征图进行归一化
        fm_s = fm_s / torch.norm(fm_s, p=2, dim=1, keepdim=True)
        fm_t = fm_t / torch.norm(fm_t, p=2, dim=1, keepdim=True)

        # 计算注意力图的差异
        loss = (fm_s - fm_t).pow(2).mean()
        return loss

# 训练循环
def train_at_distillation(teacher, student, train_loader, epochs=100):
    criterion_ce = nn.CrossEntropyLoss()
    criterion_at = ATLoss()
    optimizer = optim.Adam(student.parameters(), lr=0.001)

    for epoch in range(epochs):
        teacher.eval()
        student.train()

        for data, target in train_loader:
            data, target = data.to(device), target.to(device)

            # 获取教师和学生模型的输出及中间特征
            with torch.no_grad():
                t_output, t_features = teacher(data, return_features=True)

            s_output, s_features = student(data, return_features=True)

            # 计算损失
            loss_ce = criterion_ce(s_output, target)
            loss_at = criterion_at(s_features['layer3'], t_features['layer3'])
            total_loss = loss_ce + 0.5 * loss_at

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

我们在 CIFAR-10 数据集上进行了实验，对比了以下几种情况：