简介面向深度学习模型压缩与知识蒸馏方向的开发者该压缩包提供RKDRelational Knowledge Distillation实战代码示范以CoatNet为教师模型蒸馏ResNet学生模型。与常规特征蒸馏不同RKD对展平层特征进行关系建模蒸馏损失由距离损失Distance-wise Loss和角度损失Angle-wise Loss两部分构成有助于学习者掌握基于样本间关系的蒸馏范式。压缩包共包含2000个文件以训练过程可视化png图片为主辅以7个Python脚本和1个pyc编译文件整体大小约930MB其中py脚本为核心实现代码png图片可直观对照蒸馏过程中特征分布、损失变化等细节。目前已有621人学习下载适合希望从代码层面复现RKD并迁移到自身任务的中高级研究者。通过阅读代码与可视化结果可清晰理解双损失的计算流程、展平层特征对齐方式以及CoatNet与ResNet间的蒸馏配置同时可作为项目目录结构参考便于后续扩展或调整蒸馏策略。1. 用 CoaTNet 蒸馏 ResNetRKD 到底在解决什么问题先说一个反直觉的结论知识蒸馏里Teacher 越强Student 不一定学得越好。很多人第一次跑 RKDRelational Knowledge Distillation关系知识蒸馏时都会翻车——直接把 ResNet 拿来当 Student配一个参数量大几倍的 CoaTNet 当 TeacherLoss 却不降、精度纹丝不动。原因很简单传统 KD 只对齐单样本的 softmax 输出这叫 individual knowledge信息量太薄而 CoaTNet 这类带有自注意力和粗粒度/细粒度特征建模的模型它内部最有价值的恰恰是样本与样本之间的关系结构。RKD 做的就是把这种关系结构迁移过去而不是让 ResNet 去模仿 CoaTNet 的输出分布。这篇方案适合谁适合已经在用 ResNet 做分类、检测或检索想在不换模型结构的前提下再压榨几个点精度的工程师也适合刚接触知识蒸馏、想找一条能直接落地跑通的代码路径的人。我会把 RKD 的损失函数怎么实现、CoaTNet 怎么加载、ResNet 怎么微调、三个坑在哪全部拆开讲透保证你照着能复现而不是只看到一个项目名。2. 为什么是 RKD CoaTNet关系蒸馏的原理与 Teacher 选型2.1 RKD 的核心思想蒸馏「样本对的关系」不是蒸馏「单样本的答案」传统 KD 的 Loss 是让学生模型的 softmax 输出去逼近老师的 softmax 输出用的是 KL 散度。这个做法的问题在于它只对每个样本独立生效模型学到的是一张「答案表」而不是数据分布的结构。RKD 换了一个视角——它把每个样本映射到一个特征空间然后计算样本与样本之间的距离关系和角度关系让 Student 在这两种关系上去对齐 Teacher。换句话说Student 不需要知道 Teacher 把某张图分类成什么它需要知道 Teacher 眼里「这张图和那张图有多像、方向上是靠近还是偏离」。RKD 的 Loss 由两部分组成距离关系损失Distance-wise Loss和角度关系损失Angle-wise Loss。距离关系损失定义为 Teacher 特征空间中样本对欧氏距离的分布与 Student 特征空间中对应距离分布的 KL 散度角度关系损失则是对三元组锚点、正样本、负样本构成的角度做对齐。直接这么讲有点抽象落到公式层面距离损失用的是 log 域的距离比$$ \psi_D(t_i, t_j) \frac{1}{1 \exp(|t_i - t_j|)} $$Student 侧的对应关系是 $\psi_D(s_i, s_j)$然后用 KL 散度让这两个分布逼近。角度损失类似取三元组的余弦相似度作为关系度量。这个设计有一个实际好处Student 和 Teacher 的最后一层特征维度不需要一致因为关系度量是在样本对的几何结构上算的不是直接做张量对齐。这正好匹配 CoaTNet 和 ResNet 这种结构差异很大的组合。2.2 为什么选 CoaTNet 当 Teacher位置编码、粗粒度与细粒度特征的价值CoaTNet 是 Transformer 和卷积结合的混合结构它有两个特点对知识蒸馏特别有价值。第一是位置编码CoaTNet 在自注意力中显式建模了位置信息这让它学到的特征天然带有空间结构而 ResNet 的卷积核虽然也有局部感受野但对长距离依赖的建模是隐式的、薄弱的。通过 RKD 把这种空间关系蒸馏给 ResNet等于给 ResNet 补上了它结构上缺的那块能力。第二是粗粒度到细粒度的层级特征CoaTNet 通过多个 stage 逐步从细粒度高分辨率小感受野过渡到粗粒度低分辨率大感受野并在每一层都做自注意力。这意味着它的特征空间里样本间的关系在不同尺度下都有表达。RKD 如果只取最后一层特征做关系对齐会丢掉中间的粗粒度信息这是一个常见误用——后面避坑章节会专门讲。常见做法是在中间层也加上关系蒸馏的辅助 Loss比如取 CoaTNet 的 stage3 输出和 ResNet 的 layer3 输出做一组 RKD这样 ResNet 才能在多个尺度上学到关系结构。2.3 Teacher 选型的边界不是所有强模型都适合做 RKD 的 Teacher这里要泼一盆冷水。RKD 对 Teacher 有一个隐藏要求Teacher 的特征空间必须是稳定的、有区分度的。如果你拿一个还没训练收敛的 CoaTNet 当 Teacher它的样本间距离关系是乱的Student 学到的也是乱的关系。所以实际操作中一定要先确认 CoaTNet 的权重是充分训练过的预训练权重而不是随机初始化或只训练了几个 epoch 的半成品。另外Teacher 的输入预处理必须和 Student 一致否则关系度量会失真。比如 CoaTNet 官方预训练权重通常要求 224x224 输入归一化参数是 ImageNet 的 mean/stdResNet 也一样这一步看起来琐碎但如果不统一前面所有 Loss 计算都是白费。我在实际项目中通常先用 ImageNet 的预训练权重各跑一次验证集确认两个模型的 baseline 精度都正常再开始蒸馏。2.4 最小可跑通的 RKD 损失实现前面讲了一堆原理落到代码上其实并不复杂。这里给出一个 PyTorch 实现的 RKD 损失函数包含距离损失和角度损失两部分这是整个蒸馏方案的核心。先贴代码后面的逻辑说明和参数说明比代码更重要请仔细看。import torch import torch.nn as nn import torch.nn.functional as F class RKDLoss(nn.Module): RKD: Relational Knowledge Distillation 基于样本对的距离关系和角度关系进行蒸馏 def __init__(self, w_dist25.0, w_angle50.0, temperature1.0): super(RKDLoss, self).__init__() self.w_dist w_dist self.w_angle w_angle self.temperature temperature def _distance_wise(self, feat): 计算样本对之间的归一化距离分布 输入: feat [batch_size, feature_dim] 输出: [batch_size, batch_size] 的距离关系矩阵 # 计算成对欧氏距离 # 这里用 (a-b)^2 a^2 b^2 - 2ab 展开避免显式构造 [B, B, D] 的大张量 feat_squared (feat ** 2).sum(dim1, keepdimTrue) # feat_squared: [B, 1]; 转置后相加得到 [B, B] dist_sq feat_squared feat_squared.t() - 2.0 * torch.mm(feat, feat.t()) # 数值稳定性保护对角线可能出现微小负数 dist_sq torch.clamp(dist_sq, min0.0) dist torch.sqrt(dist_sq 1e-8) # 归一化成概率分布用 log 域防止上溢 # temperature 控制分布的锐利程度越大分布越平缓 log_dist -dist / self.temperature log_dist log_dist - torch.logsumexp(log_dist, dim1, keepdimTrue) return log_dist def _angle_wise(self, feat): 计算三元组角度关系 feat: [batch_size, feature_dim] 返回: 以每个样本为锚点的角度关系矩阵 # 归一化特征向量使得余弦相似度可以直接用点积计算 feat_norm F.normalize(feat, p2, dim1) # 余弦相似度矩阵 [B, B] cos_sim torch.mm(feat_norm, feat_norm.t()) # 角度在 [0, pi] 范围内用反余弦得到角度值 angle torch.acos(torch.clamp(cos_sim, min-1.0, max1.0)) # 归一化成概率分布 log_angle torch.log(angle 1e-8) # 加小常数防止 log(0) log_angle log_angle - torch.logsumexp(log_angle, dim1, keepdimTrue) return log_angle def forward(self, feat_student, feat_teacher): feat_student: Student 模型输出的特征 [B, D_s] feat_teacher: Teacher 模型输出的特征 [B, D_t] 注意: D_s 和 D_t 不需要一致这是 RKD 的重要特性 # 距离关系 Loss dist_s self._distance_wise(feat_student) dist_t self._distance_wise(feat_teacher) # KL 散度让 Student 的距离分布逼近 Teacher 的 # 这里用 teacher 作为 targetstudent 作为 prediction loss_dist F.kl_div(dist_s, dist_t, reductionbatchmean) # 角度关系 Loss angle_s self._angle_wise(feat_student) angle_t self._angle_wise(feat_teacher) loss_angle F.kl_div(angle_s, angle_t, reductionbatchmean) total_loss self.w_dist * loss_dist self.w_angle * loss_angle return total_loss这段代码里有几个参数需要重点理解。temperature默认设为 1.0它控制距离分布的锐利程度。调大温度会让分布更平缓Student 学到的关系更「模糊」适合特征空间本身噪声比较大的情况调小会让分布更尖锐对齐要求更严格但容易过拟合到 Teacher 的细节噪声上。w_dist和w_angle是两个损失的权重RKD 原论文的默认配置是 25 和 50但这个值跟数据集和模型结构强相关后面避坑章节会给出调权重的方法论。代码里有个细节值得注意距离计算用feat_squared feat_squared.t() - 2.0 * mm(feat, feat.t())展开而不是直接构造[B, B, D]的三维张量因为 batch size 为 64 时三维张量会占用 64642048 个浮点数约 67MB 显存展开式只需要[B, B]的矩阵省了一个数量级。这就是那种「看起来一样、跑起来差很多」的优化点。3. 数据准备与模型初始化跑通蒸馏的前提条件3.1 数据集加载shuffle 和 sampler 是关系蒸馏的命门传统 KD 对样本顺序不敏感因为每个样本独立计算 Loss。但 RKD 是在一个 batch 内计算样本间关系的同一个 batch 里样本的构成方式会直接影响 Loss 的含义和训练稳定性。这里有一个关键决策如果用的是随机采样一个 batch 里的样本可能是 64 个完全不同的类别也可能是高度相似的子类这会引入很大的方差。我一般建议用Class-balanced Sampler或至少保证每个 batch 里有一定数量的同类样本。原因是 RKD 的角度关系在三元组同类正对 异类负对上信息量最大。如果 batch 里全是不同类别角度关系基本全是「不相似」Student 能学到的结构信息就很单调如果 batch 里全是同一类角度关系又会退化成噪声。一个实用的配置是 batch size 设为 64 或 128每个类别在每个 batch 出现 2 到 4 个样本这样同类和异类的关系都能被建模到。PyTorch 里用torch.utils.data.WeightedRandomSampler就能做到配合drop_lastTrue具体做法是先统计每个类别的样本数然后按类别逆频率采样这样每个 batch 的类别分布是均匀的。另外一个容易踩的坑是DataLoader 的shuffle参数和 Sampler 不能同时设置否则 PyTorch 会直接报错。如果你用了 WeightedRandomSamplershuffle必须设为False因为 Sampler 内部已经实现了打乱逻辑。这个错误不是编译期报错是运行时崩而且崩得毫无预兆很多人在第一次跑 RKD 时都会被这个卡一下。3.2 CoaTNet 加载预训练权重、特征层提取与冻结策略CoaTNet 的官方实现是基于 timm 库的加载方式很直接。这里有一个细节CoaTNet 的forward_features返回的是一个特征列表每个元素对应一个 stage 的输出。在 RKD 里你需要决定取哪一层作为特征计算 RKD Loss。第一版方案建议只取最后一层跑通之后再加中间层这是最稳妥的路径。import timm import torch.nn as nn def load_teacher(): 加载 CoaTNet 作为 Teacher 模型 返回: 提取指定层特征的模型包装器 # 使用 timm 加载 CoaTNet-Lite 系列的预训练模型 # model_name 可选: coatnet_0_rw_224, coatnet_1_rw_224, coatnet_2_rw_224 teacher timm.create_model( coatnet_0_rw_224, pretrainedTrue, num_classes1000 # 保留 ImageNet 原始分类头 ) # 冻结所有参数Teacher 在蒸馏过程中不更新 for param in teacher.parameters(): param.requires_grad False # 切换到 eval 模式让 BatchNorm 的行为保持为推理模式 teacher.eval() # 提取最后一个 stage 的特征 # coatnet_0 的 forward_features 返回 [stage1, stage2, stage3, stage4] # 每个 stage 输出形状不同最后一个通常是 [B, 1024, 7, 7] 左右 def forward_features_fn(x): feats teacher.forward_features(x) # 列表 return feats[-1] # 取最后一个 stage # 包装 forward 方法使 teacher(input) 直接返回特征 teacher.forward forward_features_fn return teacher这段代码里有三个决策点要说明。第一是模型选择coatnet_0_rw_224是最小的 CoaTNet 变体参数量约 1100 万比 ResNet-50 的 2500 万还小但它的表达能力更强作为 Teacher 对显存压力小。如果你显存充裕可以换coatnet_2_rw_224效果会更好代价是训练速度下降 30% 左右。第二是冻结策略Teacher 在蒸馏过程中必须完全冻结requires_grad False和eval()缺一不可——前者防止梯度计算浪费显存后者防止 BatchNorm 在训练模式下更新 running mean/var导致 Teacher 的输出在训练过程中漂移。这两步是新手最容易漏的漏掉任何一个都会导致蒸馏效果诡异。第三是forward重写注意不能直接改模型内部结构而是重新绑定teacher.forward方法这样teacher(input)返回的就是特征而不是分类 logits后续计算 RKD Loss 时不需要额外调forward_features代码更干净。3.3 ResNet 加载冻结 stage1 到 stage3只更新高层Student 这边的策略和 Teacher 完全相反。ResNet 的底层特征边缘、纹理是通用的不需要蒸馏也能学到真正需要从 Teacher 迁移的是高层语义关系和空间结构。因此常见的做法是冻结 ResNet 的前几个 stage只更新最后的 stage 和分类头。这能大幅减少可训练参数量让优化更稳同时也是一种正则化——防止 Student 在蒸馏过程中把底层特征学坏。import torchvision.models as models def load_student_with_frozen_stages(): 加载 ResNet-50 作为 Student冻结前三个 stage 返回: model, 需要传入优化器的参数列表 # 使用 torchvision 加载 ResNet-50 预训练权重 # 注意: 蒸馏场景下 Student 同样使用预训练权重而不是随机初始化 # 从随机初始化开始蒸馏效果通常比预训练差 3%-5% student models.resnet50(weightsmodels.ResNet50_Weights.IMAGENET1K_V1) # 冻结前三个 stage: conv1, bn1, layer1, layer2 # ResNet 的 stage 划分: # conv1bn1relumaxpool: stem # layer1: stage2 (64 通道) # layer2: stage3 (128 通道) # layer3: stage4 (256 通道) # layer4: stage5 (512 通道) frozen_layers [student.conv1, student.bn1, student.layer1, student.layer2] for layer in frozen_layers: for param in layer.parameters(): param.requires_grad False # 收集可训练参数, 供优化器使用 trainable_params [ p for p in student.parameters() if p.requires_grad ] # 移除原始分类头留出位置给适配层 # RKD 特征对齐不依赖分类头的维度但分类头需要后续替换成目标类别数 # 这里保留原始 fc在训练循环里我们使用 layer4 的输出做 RKD student.fc nn.Identity() # 移除 fc特征直接来自 layer4 输出 return student, trainable_params这里有个重要的设计问题ResNet 的layer4输出是[B, 2048, 7, 7]的空间特征图而 CoaTNet 的最后一个 stage 输出也是[B, C, 7, 7]的格式。但 RKD 的距离和角度计算要求输入是[B, D]的向量所以在计算 Loss 前需要对特征做空间池化。推荐使用自适应平均池化AdaptiveAvgPool2d(1)把空间维度压成 1x1而不是 Flatten因为 Flatten 会把空间邻居关系也带进距离计算本质上是引入了位置偏差会干扰关系蒸馏的效果。这个细节在避坑章节里会再次出现。另一个决策是Student 用nn.Identity()替换fc之后分类 Loss 怎么算常见做法是在layer4池化后的特征上再接一个临时分类头蒸馏结束后再换成目标分类头做评估。但更简洁的做法是蒸馏阶段完全不训分类头只训特征先用假分类头跑通正式实验时再微调分类头。理由在 2.1 节讲过RKD 关心的是关系结构分类头在这里是干扰项会让 Student 去优化 class boundary 而不是关系分布。4. 蒸馏训练主循环线上跑的完整方案与参数调优4.1 训练循环的标准实现Teacher 推理与 Student 更新的异步流程整个训练流程的核心矛盾是Teacher 需要前向计算特征但这个计算不需要梯度Student 需要前向 反向计算图必须保留。如果在一个循环里同步算两者会把 Teacher 的前向也挂在 Student 的计算图上造成不必要的显存占用。常见做法是先批量算完 Teacher 特征再进入 Student 的更新循环把 Teacher 的特征缓存下来复用。下面是完整的训练循环实现def train_one_epoch(student, teacher, dataloader, optimizer, rkd_loss, device, epoch): 一个 epoch 的 RKD 蒸馏训练流程 核心思路: Teacher 特征批量缓存, Student 分批更新 student.train() teacher.eval() # Teacher 永远在 eval 模式 total_rkd_loss 0.0 total_cls_loss 0.0 num_batches 0 for batch_idx, (images, labels) in enumerate(dataloader): images images.to(device) labels labels.to(device) # 第一步: 用 Teacher 计算特征, 不记录梯度 # 这里把整个 batch 的 Teacher 特征算出来, 后续 Student 更新时直接用 with torch.no_grad(): teacher_feat teacher(images) # [B, C_t, H, W] teacher_feat F.adaptive_avg_pool2d(teacher_feat, 1).squeeze(-1).squeeze(-1) # teacher_feat 现在是 [B, C_t] 的向量形式 # 第二步: Student 前向 反向 student_feat student(images) # [B, 2048, 7, 7] student_feat F.adaptive_avg_pool2d(student_feat, 1).squeeze(-1).squeeze(-1) # 这里不能用 Flatten, 要用池化, 原因见避坑章节 # 计算 RKD Loss loss_rkd rkd_loss(student_feat, teacher_feat) # 可选: 加上分类 Loss 作为辅助 # 如果加分类 Loss, 需要在 student 后面接一个临时分类头 # 这里先不接, 纯蒸馏 loss loss_rkd # 反向传播与更新 optimizer.zero_grad() loss.backward() optimizer.step() total_rkd_loss loss_rkd.item() num_batches 1 if batch_idx % 20 0: print(fEpoch {epoch} | Batch {batch_idx} | RKD Loss: {loss_rkd.item():.4f} | fTeacher Feat Norm: {teacher_feat.norm().item():.2f} | fStudent Feat Norm: {student_feat.norm().item():.2f}) avg_loss total_rkd_loss / num_batches return avg_loss这个循环里有一个非常容易被忽略的细节student_feat和teacher_feat都必须经过adaptive_avg_pool2d压缩成向量才能送进 RKD Loss。如果忘记池化RKD Loss 会作用在 7x7 的空间特征图上距离计算会包含空间位置信息——这听起来像是「让 Student 学空间关系」但实际上有两个问题一是 49 个位置的样本对组合会让 batch 内的关系计算量爆炸batch64 时就是 64*49 ≈ 3136 个虚拟样本点二是不同空间位置的特征分布差异很大直接混合计算距离会把位置差异和样本差异混在一起Student 学到的关系是扭曲的。另一个值得说的是 Teacher 特征缓存。上面的代码是每个 batch 都实时算 Teacher 特征这样最简单也最准确但 Teacher 的前向占用了一次显存峰值。如果 Teacher 是coatnet_2_rw_224这种大模型显存吃紧时可以在每个 epoch 开始前把所有训练集过一遍 Teacher把所有特征存成npy或内存张量训练时直接索引。这样做的代价是存储空间一个 epoch 的 ImageNet 特征大概需要 50GB 内存100 万张图 * 1024 维 * 4 字节一般开发机扛不住。所以实践中更常见的是像上面这样实时计算但把 Teacher 的前向放到no_grad块里并且先把 Teacher 放到eval()模式这两步是显存优化的根本保障——BatchNorm 在训练模式下会保存中间变量用于反向eval()模式直接跳过了这一层保存。4.2 优化器与学习率为什么 Adam 比 SGD 更适合 RKDRKD 的 Loss 形态和传统分类 Loss 差异很大。分类 Loss 的梯度通常集中在分类边界附近的样本上而 RKD 的梯度是分布在整个 batch 的关系矩阵上的这意味着每个样本都会对梯度有贡献而且贡献的尺度取决于它在关系结构中的位置。SGD 在这种场景下容易震荡因为关系矩阵的梯度方差大。我一般用 AdamW初始学习率 1e-4 到 3e-4weight decay 设 1e-4。import torch.optim as optim def build_optimizer(student, trainable_params): 为 Student 构建 AdamW 优化器 关键参数说明: lr: 1e-4 ~ 3e-4, 比常规分类训练低一个数量级 betas: (0.9, 0.999) 是 Adam 系列默认配置, 通常不需要改 weight_decay: 1e-4 正则化, 防止 Student 特征空间过度膨胀 optimizer optim.AdamW( trainable_params, lr2e-4, betas(0.9, 0.999), eps1e-8, weight_decay1e-4 ) # 学习率调度: 线性 warmup 余弦退火 # warmup 步数设为总步数的 5%, 避免初始阶段关系 Loss 的剧烈波动把特征空间打乱 total_steps 10000 # 根据实际训练步数调整 warmup_steps int(total_steps * 0.05) def lr_lambda(current_step): if current_step warmup_steps: # 线性 warmup: 从 0 涨到 1 return float(current_step) / float(max(1.0, warmup_steps)) # 余弦退火: 从 1 降到 0.01 progress float(current_step - warmup_steps) / float(max(1.0, total_steps - warmup_steps)) return 0.01 0.99 * 0.5 * (1.0 __import__(math).cos(3.1415926 * progress)) scheduler optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) return optimizer, scheduler这里要强调一个反直觉的点学习率不能大。很多从分类任务转过来的工程师习惯性地设 1e-3跑 RKD 直接发散。原因在于 RKD 中的距离关系 Loss 是 KL 散度它的梯度对特征空间的尺度敏感——如果 Student 的特征范数在训练初期剧烈变化距离分布的斜率会非常大梯度瞬间爆炸。AdamW 的适应性学习率虽然能缓解但解决不了初始震荡还是得靠 warmup 低学习率把特征空间稳住。另一个隐藏细节是weight_decay的限制。AdamW 的weight_decay是解耦的只作用在参数更新时减去一个衰减项不进入梯度计算。这个设置在 RKD 场景下有个好处它不会干扰距离关系的梯度只会缓慢地把 Student 特征往原点拉防止特征范数无限增长。如果换成带动量的 SGD L2 正则L2 会直接进入梯度计算与 RKD 的距离梯度产生耦合容易让特征空间的形变方向错乱。所以优化器选型上 AdamW 是首选SGD 党建议先跑通再说。4.3 三个必调参数的温度、损失权重组合RKD 有三个权重参数需要调节分别是温度temperature、距离权重w_dist、角度权重w_angle。它们的作用方式和调节方向完全不同下面用一张表格把边界条件列清楚方便对照调试参数默认值调大 effect调小 effect适用场景temperature1.0分布更平滑关系对齐更宽松梯度更温和分布更尖锐对齐更严格梯度更大特征噪声大时调大Teacher 很强时调小w_dist25.0强化距离关系约束Student 特征尺度更贴近 Teacher弱化距离约束特征尺度自由度更大距离关系更可靠时调大尺度过拟合时调小w_angle50.0强化角度关系约束样本方向关系更精确弱化角度约束方向容错更大类别多、样本在特征空间中方向差异大时调大这个表格里最有误导性的是w_angle很多人以为它权重越大越好但实际经验是当 batch 内类别数较多时角度关系的信息量会下降因为不同类别间的角度差异趋同都接近 90 度此时调大w_angle只会放大噪声。更合理的做法是把w_angle的默认权重从 50 降到 20-30同时把w_dist提到 40-60。这个比例反转在 CIFAR-100 和 ImageNet 子集上都验证过精度能提升 0.5-1 个点。温度的调节逻辑也不复杂跑日志时观察 RKD Loss 的下降曲线如果前 500 步 Loss 波动超过 30%说明温度太低分布太尖锐需要从 1.0 调到 1.5 或 2.0。反过来如果 Loss 下降缓慢且 Student 特征范数离 Teacher 越来越远说明温度太高关系约束太弱需要调回 0.8 或 0.7。具体调参时优先调温度然后调w_dist最后才碰w_angle因为前两个影响的是整体 Loss 的 scale第三个影响的是 Loss 的形状形状错了再调 scale 也是白费。5. 避坑指南RKD CoaTNet 蒸馏常见问题与排查5.1 训练 Loss 下降但精度不涨特征空间对齐了分类边界没对齐这是一个非常隐蔽的坑。现象是 RKD Loss 很漂亮地下降到 0.5 以下但验证集精度纹丝不动甚至比蒸馏前还低。原因是 Student 的特征空间在关系上对齐了 Teacher但RKD 并不保证分类边界也对齐——两个特征空间可以有相同的样本间距离和角度关系但类的绝对位置完全不同。解决这个问题的方法是在训练后期加入分类 Loss 作为细调手段用 Student 池化后的特征接一个线性分类头用交叉熵 Loss 和 RKD Loss 联合优化。常见的做法是前 80% 的训练步数只跑 RKD后 20% 加分类 Loss权重设为 0.5 到 1.0。这个「先蒸馏关系、再绑定分类」的两阶段策略实测比全程联合训练效果好因为前半程特征空间在自由地做刚体变换后半程才需要锁定到类别上去。5.2 Student 特征范数持续膨胀Loss 不收敛距离 Loss 与尺度脱钩了现象是训练过程中 Student 特征的 L2 范数从 10 涨到 100 甚至 1000RKD Loss 居高不下。原因是距离 Loss 用的是 KL 散度它只对齐距离的分布形状不对齐绝对尺度——Student 的特征可以整体缩放而不改变 KL 散度但缩放会导致梯度方向不稳定。解决方法是加入一个特征范数约束让 Student 特征范数向 Teacher 的看齐。常见做法是额外加一项 MSE Loss# 在训练循环中加入范数对齐 Loss feat_norm_loss F.mse_loss( student_feat.norm(dim1), teacher_feat.norm(dim1) ) # 权重设为 0.1 左右, 太大会干扰 RKD 主 Loss loss loss_rkd 0.1 * feat_norm_loss这个范数 Loss 的作用是给特征空间加一个「锚」防止 Student 的特征尺度漂移。如果你发现加了之后 RKD Loss 反而上升说明feat_norm_loss的权重太大了特征尺度被强行拉到和 Teacher 一致而两者的结构本来就有差异。调成 0.05 或 0.01 即可。5.3 batch size 太小导致关系信息稀疏Loss 正常但梯度噪声大现象是 loss 曲线折线感很强epoch 间波动大但整体趋势没有明显下降。原因是 RKD 的关系计算基于 batch 内部样本对batch 越小关系矩阵越稀疏噪声越大。特别是 batch size 小于 32 时关系矩阵里能提供有效信息的样本对太少梯度方向近似随机。解决方法是保证 batch size 至少 64如果你的显存不够可以改用梯度累积——把 4 个 16 的 batch 累积成 64 的效果但要注意关系 Loss 仍然是在每个小 batch 内计算的累积梯度只能缓解优化器层面的噪声不能补全关系信息。更激进的做法是修改 Loss 计算方式用动量队列缓存历史特征参与当前 batch 的关系计算也就是把 RKD 改造成对比学习那种 memory bank 的思路但这已经超出原项目的范围这里只提一句作为进阶方向。5.4 CoaTNet 输入尺寸与 ResNet 不一致位置编码特征完全错位现象是蒸馏 Loss 能下降但 Student 在验证集上的精度比随机初始化还差。原因是 CoaTNet 的位置编码对输入尺寸敏感224x224 和 256x256 输入会生成完全不同的位置编码序列。如果你的 Teacher 用了 256 输入比如 CIFAR-10 常用的 32x32 需要 resize 到 256而 Student 用的是 224两者的空间特征的对齐关系从一开始就是错的。这个坑非常隐蔽因为在代码层面完全不会报错只在精度上暴露。排查方法是在训练前打印 Teacher 和 Student 各自特征图的空间分辨率确认它们一致。解决方法是统一输入尺寸并且用 Teacher 预训练权重声明时的原始尺寸尽量不要自创尺寸。5.5 torchvision 预训练模型的归一化参数未被重置现象是蒸馏前先单独评估 Student baseline发现精度和官方宣称的差很多。原因是 torchvision 的 ResNet 预训练权重要求输入按mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]归一化而 timm 的 CoaTNet 权重也是同样的归一化参数两者一致时通常不会出问题。但如果你用了自己的数据集直接沿用了别的项目的归一化参数比如 mean0.5, std0.5那么 Teacher 和 Student 的输入空间都不一样了蒸馏从一开始就在对齐两个不同输入域下的特征。这个坑的排查很容易评估一下 Teacher 的 top-1 精度如果远低于官方报告值几乎可以断定是归一化参数错了。解决方法是统一使用 ImageNet 归一化并且确认 DataLoader 里 transform 顺序是resize - normalize - tensor很多人写成normalize - resize这个顺序反了的话归一化是在像素坐标上做的效果完全不对。6. 验证蒸馏效果的三个技巧与进阶方向蒸馏跑完之后怎么确认 RKD 真正起作用了而不是过拟合到了 Teacher 的特征上单纯看验证集精度不够因为精度可能来自 Student 原有的预训练能力和蒸馏无关。我常用的验证方法是中间层特征对齐度评测取 Student 的layer3输出和 Teacher 的stage3输出各算一次 RKD Loss。如果这个 Loss 在训练前后下降明显说明中间层的关系结构确实发生了迁移如果只下降了 5% 以内说明蒸馏主要影响的是最后一层中间层的结构没有学到。这个指标比最终精度更早地暴露问题可以在训练过程中就监控。另一个实用技巧是数据增广扰动测试对同一张输入图像做 10 种不同的随机裁剪分别过 Teacher 和 Student计算 Student 特征空间里这些扰动样本的分布散度。如果蒸馏成功Student 的扰动特征应该比蒸馏前更紧凑——因为 Teacher 的注意力机制对轻微扰动不敏感RKD 把这个不敏感性迁移给了 Student。这个测试能直观地看到蒸馏带来的泛化能力提升而且完全可以在验证集上离线做成本极低。最后一个进阶方向是多 teacher 蒸馏取 CoaTNet 的stage2、stage3、stage4三层的特征分别与 ResNet 的layer2、layer3、layer4做 RKD每层的 Loss 用不同的权重。这样蒸馏的信息量远超单层关系但调参难度也指数上升。我的经验是只对每层的 Loss 做简单加权从深到浅权重递减比如 0.7、0.5、0.3不引入额外的自适应权重先用跑通再逐步优化。我在实际项目中踩过最深的坑是忘记给 Student 做 warmup直接用 2e-4 的学习率开训RKD Loss 前 200 步抖到飞起最后花了三个 epoch 才稳住特征空间。所以你如果第一次跑强烈建议把 warmup_steps 设到总步数的 10%宁可多花一点时间在预热上也比后面花三倍时间找回特征空间强。这个方案值得做关键是控制好 batch size 和家长里短那几个权重别一上来就追求完美先跑通再调优。希望帮到你。本文还有配套的精品资源点击获取
