基于注意力机制的特征蒸馏（AFD）实战指南：从论文解读到模型优化

1次阅读

没有评论

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

在模型压缩和迁移学习中，特征蒸馏是一种常用技术，目的是让小模型（学生模型）学习大模型（教师模型）的特征表示能力。然而传统方法存在明显痛点：

跨架构迁移困难：当教师模型和学生模型结构差异较大时（如 ResNet 到 MobileNet），直接使用 L2 损失强制对齐特征图会导致信息损失
特征空间不对齐：不同网络层的特征响应分布差异大，简单的最小二乘匹配会强制学生模型模仿无关噪声
忽略关键区域：传统方法平等对待所有空间位置，无法聚焦于教师模型真正 ” 关注 ” 的重要特征区域

AFD 的核心创新是通过注意力机制动态调整特征蒸馏的权重分配，其流程分为三个阶段：

通道注意力：识别重要特征通道
$$W_c = \sigma(MLP(AvgPool(F_T)))$$
其中 $F_T$ 是教师模型特征图，$\sigma$ 表示 sigmoid 函数
空间注意力：定位关键空间区域
$$W_s = \sigma(conv^2(concat[AvgPool(F_T), MaxPool(F_T)]))$$
实例注意力：样本自适应加权
$$W_i = \frac{exp(S_T/\tau)}{\sum exp(S_T/\tau)}$$
其中 $S_T$ 是教师模型的预测置信度，$\tau$ 为温度系数

对比原论文在 CIFAR-100 上的实验数据：

方法	学生准确率	提升幅度
L2 蒸馏	72.1%	–
KL 蒸馏	73.8%	+1.7%
AFD(本文)	76.3%	+4.2%

以下是完整的 AFD 损失函数实现：

import torch
import torch.nn as nn
import torch.nn.functional as F

class AFDLoss(nn.Module):
    def __init__(self, temp=4.0):
        super().__init__()
        self.temp = temp
        # 通道注意力 MLP
        self.mlp = nn.Sequential(nn.Linear(512, 32),  # 输入维度需适配
            nn.ReLU(),
            nn.Linear(32, 512)
        )
        # 空间注意力卷积
        self.conv = nn.Conv2d(2, 1, kernel_size=7, padding=3)

    def forward(self, feat_s, feat_t, logits_t):
        # 通道注意力
        b, c, _, _ = feat_t.shape
        channel_att = torch.sigmoid(self.mlp(F.avg_pool2d(feat_t, feat_t.size()[2:]).view(b, c)))

        # 空间注意力
        avg_pool = torch.mean(feat_t, dim=1, keepdim=True)
        max_pool, _ = torch.max(feat_t, dim=1, keepdim=True)
        concat = torch.cat([avg_pool, max_pool], dim=1)
        spatial_att = torch.sigmoid(self.conv(concat))

        # 实例注意力
        instance_att = F.softmax(logits_t/self.temp, dim=1)

        # 组合注意力
        attention_map = channel_att.view(b,c,1,1) * spatial_att * instance_att.view(b,1,1,1)

        # 加权特征差异
        loss = (attention_map * (feat_s - feat_t)**2).sum() / b
        return loss

关键实现细节：
– 特征图维度需要根据实际模型调整
– 温度系数默认设为 4.0，可根据任务调节
– 三种注意力采用逐元素相乘方式组合

分类任务：
温度系数 τ：一般 3.0-5.0，简单任务取低值
损失权重 λ：建议 0.5-1.0，与交叉熵损失配合使用
检测任务：
需对 ROI 特征单独计算 AFD 损失
建议冻结骨干网络前几层不参与蒸馏

使用梯度缩放 (gradient scaling) 避免数值下溢

loss = afd_loss(feat_s, feat_t, logits_t) * batch_size / 256

混合精度训练时注意对 attention_map 做 float32 转换

结合 Grad-CAM 观察注意力效果：

# 以 ResNet 为例
target_layer = model.layer4[-1]
cam = GradCAM(model, target_layer)
grayscale_cam = cam(input_tensor)
# 叠加原始图像
visualization = show_cam_on_image(img, grayscale_cam)