4. Feature蒸馏实战

Feature蒸馏，说白了就是让学生的中间层特征图去模仿老师的。我刚开始做蒸馏时，只盯着logits做KL散度，效果确实有提升，但总觉得学生学得不够「透」。后来试了试中间层对齐，才发现——原来老师中间层的特征图里，藏着那么多学生没学到的东西。

这一章，我们重点聊三种主流方法：FitNet、Attention Transfer（AT）和Similarity Preserving（SP）。我会把原理、代码、还有我踩过的坑，一次性说清楚。

4.1 FitNet：直接回归特征图

FitNet的思路很直接——把学生某层的特征图，通过一个适配层（通常是1x1卷积或全连接），强行对齐到老师对应层的特征图。损失函数就是简单的MSE。

核心公式：

L_FitNet = || f_t(x) - f_s(x) ||²

其中 f_t 是老师特征，f_s 是学生特征（经过适配层后）。

代码实现其实不复杂。我贴一个我项目里用过的版本：

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

class FitNetLoss(nn.Module):
    def __init__(self, student_channels, teacher_channels, hidden_dim=256):
        super().__init__()
        # 适配层：把学生特征通道数映射到老师通道数
        self.adaptor = nn.Sequential(
            nn.Conv2d(student_channels, hidden_dim, kernel_size=1),
            nn.BatchNorm2d(hidden_dim),
            nn.ReLU(),
            nn.Conv2d(hidden_dim, teacher_channels, kernel_size=1)
        )
    
    def forward(self, student_feat, teacher_feat):
        # student_feat: [B, C_s, H, W]
        # teacher_feat: [B, C_t, H, W]
        adapted = self.adaptor(student_feat)
        # 如果空间尺寸不一致，需要插值
        if adapted.shape[2:] != teacher_feat.shape[2:]:
            adapted = F.interpolate(adapted, size=teacher_feat.shape[2:], mode='bilinear')
        loss = F.mse_loss(adapted, teacher_feat.detach())
        return loss

注意：老师特征一定要detach()，否则梯度会回传到老师网络。我刚开始写的时候忘了这步，结果老师参数也跟着更新了，训练直接崩了。

4.2 Attention Transfer（AT）：关注激活模式

直接回归特征图有个问题——老师特征图的每个像素不一定都重要。AT的思路是：我们只关心特征图的「注意力分布」，也就是哪些区域被激活得更强。

AT的损失函数是这样：

L_AT = || A_t / ||A_t||₂ - A_s / ||A_s||₂ ||²

其中A是注意力图，通常用特征图在通道维度的平方和（sum of squares）来计算。

我习惯这么写：

class ATLoss(nn.Module):
    def __init__(self, p=2):
        super().__init__()
        self.p = p  # 幂次，通常用2
    
    def attention_map(self, feat):
        # feat: [B, C, H, W]
        # 计算每个空间位置的激活强度
        attn = torch.sum(feat ** self.p, dim=1)  # [B, H, W]
        # 归一化
        norm = torch.norm(attn.view(attn.size(0), -1), dim=1, keepdim=True)
        attn = attn / (norm.view(-1, 1, 1) + 1e-8)
        return attn
    
    def forward(self, student_feat, teacher_feat):
        s_attn = self.attention_map(student_feat)
        t_attn = self.attention_map(teacher_feat)
        loss = F.mse_loss(s_attn, t_attn.detach())
        return loss

AT的好处是——它不要求学生和老师的特征图通道数一致，甚至空间尺寸不同也能用（大不了插值一下注意力图）。我在做轻量化网络蒸馏时，经常用AT配合logit蒸馏，效果比单独用FitNet稳定很多。

4.3 Similarity Preserving（SP）：保持样本间相似性

SP的思路又不一样了。它不直接对齐特征图，而是对齐「样本之间的相似度矩阵」。说白了，就是让学生的特征空间中，样本A和样本B的距离关系，跟老师特征空间中的距离关系保持一致。

损失函数长这样：

L_SP = || G_t / ||G_t||_F - G_s / ||G_s||_F ||²

其中G是Gram矩阵（样本间的内积矩阵），||·||_F是Frobenius范数。

代码实现：

class SPLoss(nn.Module):
    def __init__(self):
        super().__init__()
    
    def forward(self, student_feat, teacher_feat):
        # 展平特征图
        s_flat = student_feat.view(student_feat.size(0), -1)  # [B, C*H*W]
        t_flat = teacher_feat.view(teacher_feat.size(0), -1)
        
        # 计算Gram矩阵（样本间相似度）
        s_gram = torch.mm(s_flat, s_flat.t())  # [B, B]
        t_gram = torch.mm(t_flat, t_flat.t())
        
        # 归一化
        s_gram = s_gram / torch.norm(s_gram, p='fro')
        t_gram = t_gram / torch.norm(t_gram, p='fro')
        
        loss = F.mse_loss(s_gram, t_gram.detach())
        return loss

我的经验：SP在batch size较小时效果会打折扣。因为Gram矩阵是B×B的，batch太小，样本间的关系就不够丰富。我一般建议batch size至少32以上。

4.4 三种方法的对比

方法 对齐目标 通道数要求 空间尺寸要求 适用场景
FitNet 特征图直接回归 需要适配层 需要插值对齐 学生容量较大时
AT 注意力分布 无要求 可插值对齐 通用，稳定性好
SP 样本间相似度 无要求 无要求 batch size较大时

4.5 中间层特征对齐的坑与技巧

这部分是我最想说的。Feature蒸馏看着简单，但实际落地时坑特别多。我一个个讲。

坑1：层选择不对，效果反而变差

不是所有层都适合做蒸馏。我试过把学生第一层和老师最后一层对齐，结果学生直接学废了。一般来说，选择网络中间偏后的层效果最好。我个人习惯选总层数的后1/3到1/2位置。

坑2：损失权重配比

Feature蒸馏的损失通常要和logit蒸馏、hard loss一起用。权重怎么配？我踩过很多次坑。一个比较稳的起点是：hard loss权重1.0，logit蒸馏权重0.5~1.0，feature蒸馏权重0.1~0.5。具体还得调，但至少不会崩。

坑3：适配层设计

FitNet的适配层不是越深越好。我试过用3层卷积做适配，结果学生反而学不到东西。后来发现，1层1x1卷积就够了，最多加个BN和ReLU。适配层太深，梯度信号会被稀释。

坑4：特征图尺寸不一致

老师和学生的特征图空间尺寸经常不一样。直接用插值？可以，但要注意——插值会引入额外的平滑效应。我建议尽量选择尺寸一致的层做对齐，实在不行再用插值。

我曾经踩过的坑：有一次我用AT蒸馏，注意力图归一化时忘了加epsilon（1e-8），结果某个batch的特征图全零，除以0直接NaN了。从那以后，我所有归一化操作都加epsilon。

4.6 知识体系总览

下面这张图，把三种方法的核心逻辑和关系画清楚了：

Feature蒸馏三大方法对比 Teacher Network Student Network Teacher Feature Student Feature FitNet: 直接MSE回归 AT: 注意力分布对齐 SP: 样本相似性保持 选择建议 • 学生容量大 → FitNet（直接学特征细节） • 通用场景 → AT（稳定，对通道数不敏感） • batch size大且关注关系 → SP（保持样本间结构）

4.7 实战建议

最后，我总结几条实战经验：

  • 先跑通logit蒸馏，再加feature蒸馏。我习惯先让学生学会模仿老师的输出分布，再用feature蒸馏做精细化调整。
  • Feature蒸馏的权重从0.1开始调。权重太大，学生容易过拟合到老师的特征细节上，反而泛化变差。
  • 多尝试组合使用。比如AT+SP一起用，效果往往比单独用任何一种都好。但要注意总损失不要太大。
  • 监控特征图的可视化。我每次训练都会随机抽几张图，看看学生和老师的注意力图是否相似。如果差异很大，说明蒸馏还没到位。

一个小技巧：如果你不确定选哪一层做蒸馏，可以先把老师所有层的特征图都存下来，然后逐层计算学生和老师的相似度（用余弦相似度或CKA）。选相似度最低的那一层做蒸馏，效果往往最好。

Feature蒸馏这块，方法看着多，但核心就一句话——让学生学会老师「怎么看」和「怎么想」。FitNet学的是「看什么」，AT学的是「关注哪里」，SP学的是「怎么关联」。理解了这个，代码怎么写都不难。

← 上一章 📖 返回目录 下一章 →