基于AFD注意力特征蒸馏的模型压缩实战：从论文到工业部署

1次阅读

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

在移动设备和边缘计算场景中，深度神经网络模型常常面临计算资源有限、功耗受限的挑战。传统的知识蒸馏（Knowledge Distillation, KD）通过让小型学生模型模仿大型教师模型的输出，已经证明是一种有效的模型压缩方法。然而，标准的 KD 方法存在一个关键限制——它主要关注最终预测层的对齐，而忽略了中间层特征的丰富信息。

这就是 AFD（Attention-based Feature Distillation）注意力特征蒸馏技术的用武之地。AFD 通过引入注意力机制，让学生模型能够更有针对性地学习教师模型的关键特征表示，从而在保持精度的同时实现更高效的压缩。

传统 KD 和 AFD 都遵循 ” 教师 - 学生 ” 的学习框架，但它们在特征转移的方式上存在根本差异：

传统 KD 仅使用教师模型的 softmax 输出作为监督信号
AFD 则利用注意力机制提取教师模型中间层的空间和通道特征重要性
AFD 可以自动识别并对齐教师和学生模型中最有价值的特征区域

这种差异带来的实际效果是：

AFD 能够保留更多细粒度的特征信息
注意力机制使学生模型能更高效地利用有限的参数容量
在相同压缩率下，AFD 通常能获得更高的精度

AFD 的核心创新在于它的注意力机制设计。具体来说，它包含两种注意力：

空间注意力：识别特征图中哪些空间位置更重要
通道注意力：识别哪些特征通道更具代表性

这些注意力权重是通过简单的全连接层和小型神经网络计算得到的，计算过程如下：

# 空间注意力计算示例
class SpatialAttention(nn.Module):
    def __init__(self, in_channels):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, 1, kernel_size=1)

    def forward(self, x):
        # x 形状: [B, C, H, W]
        att = self.conv(x)  # [B, 1, H, W]
        att = F.softmax(att.view(x.size(0), -1), dim=1).view_as(att)
        return att

AFD 的损失函数由三部分组成：

标准分类损失（学生模型输出与真实标签的交叉熵）
注意力特征蒸馏损失（教师与学生注意力权重的 KL 散度）
特征激活蒸馏损失（教师与学生特征图的 L2 距离）

数学表达为：

L_total = α*L_cls + β*L_att + γ*L_feat

其中 α、β、γ 是超参数，通常我们设置 β =γ=0.5，α=1.0 作为初始值。

下面给出 AFD 的核心实现模块：

class AFD(nn.Module):
    """
    注意力特征蒸馏模块
    参数:
        in_channels: 输入特征图的通道数
        T: 温度系数，默认为 4
    """
    def __init__(self, in_channels, T=4):
        super().__init__()
        self.T = T
        # 空间和通道注意力模块
        self.spatial_att = SpatialAttention(in_channels)
        self.channel_att = ChannelAttention(in_channels)

    def forward(self, feat_s, feat_t):
        """
        参数:
            feat_s: 学生特征图 [B, C, H, W]
            feat_t: 教师特征图 [B, C, H, W]
        返回:
            蒸馏损失值
        """
        # 计算空间注意力
        s_spatial = self.spatial_att(feat_s)
        t_spatial = self.spatial_att(feat_t)

        # 计算通道注意力
        s_channel = self.channel_att(feat_s)
        t_channel = self.channel_att(feat_t)

        # 计算 KL 散度损失
        loss_spatial = F.kl_div(F.log_softmax(s_spatial.view(-1)/self.T, dim=0),
            F.softmax(t_spatial.view(-1)/self.T, dim=0),
            reduction='batchmean') * (self.T ** 2)

        loss_channel = F.kl_div(F.log_softmax(s_channel/self.T, dim=1),
            F.softmax(t_channel/self.T, dim=1),
            reduction='batchmean') * (self.T ** 2)

        # 特征激活损失
        loss_feat = F.mse_loss(feat_s, feat_t)

        return loss_spatial + loss_channel + loss_feat

在训练过程中，我们需要将 AFD 模块集成到标准训练流程中：

def train_with_afd(student, teacher, train_loader, afd):
    student.train()
    teacher.eval()  # 教师模型不更新参数

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

        # 前向传播
        with torch.no_grad():
            feat_t, out_t = teacher(images, return_features=True)

        feat_s, out_s = student(images, return_features=True)

        # 计算各类损失
        loss_cls = F.cross_entropy(out_s, labels)
        loss_afd = afd(feat_s, feat_t)

        # 总损失
        loss = loss_cls + 0.5 * loss_afd

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

我们在 CIFAR-10 数据集上对比了以下几种方法：