简介基于RKDRelational Knowledge Distillation的知识蒸馏实战代码包面向有一定深度学习基础、想对比CoatNet与ResNet蒸馏策略的算法工程师、研究生及竞赛选手。RKD与常规特征蒸馏不同聚焦于展平层特征的关系建模不直接对齐教师与学生特征而是通过二阶距离损失Distance-wise Loss拉近样本对之间的相对距离并用三阶角度损失Angle-wise Loss保持三元组的角度结构从而提升学生网络的表征能力。包内共2000余个文件以2406张训练过程与可视化PNG图片为主可直观观察损失曲线、特征分布和蒸馏效果另有7个Python脚本和1个pyc文件覆盖模型定义、CoatNet教师与ResNet学生构建、蒸馏训练与评估等核心逻辑。压缩包整体约930.94MB目录结构清晰适合直接运行、调试和二次修改。目前已有621人学习下载是复现RKD损失细节、完成CoatNet蒸馏ResNet任务的实用参考。1. RKD知识蒸馏实战CoatNet当教师ResNet能学到什么很多人一听到用CoatNet蒸馏ResNet第一反应是方向反了CoatNet的参数量并不比ResNet-50大多少凭什么当教师这正是RKD这类知识蒸馏方法要解决的问题。RKDRelational Knowledge Distillation蒸馏的不是单个样本的输出而是样本与样本之间的关系——距离远近和夹角大小所以教师和学生可以拥有完全不同的结构。CoatNet带自注意力和位置编码它看到的特征空间里粗粒度分类边界和细粒度判别结构都比纯卷积的ResNet清晰用RKD把这种几何关系迁到ResNet的特征空间里学生不用复制CoatNet的结构也能获得更接近它的判别方式。这篇会按真实落地顺序展开先讲清RKD原理和选型理由再给出改造ResNet、接入CoatNet的框架代码然后落到训练循环、损失参数和避坑最后介绍一套更有说服力的验证方法。适合正在做分类模型压缩、想给卷积网络注入注意力机制知识的工程师阅读。2. RKD原理与教师选择为什么关系蒸馏更适合CoatNet到ResNet2.1 知识蒸馏的两种形态输出对齐与关系对齐Hinton式知识蒸馏HKD把教师模型的logits当作软目标让学生在分类任务上匹配这个概率分布。这个做法在师生结构相似时很省事实现只有几行代码。但它有一个隐含假设教师和学生最后的输出分布能够在相同语义下比较好。CoatNet和ResNet的输出分布虽然都是1000类概率可两个模型学习到的特征空间结构差得很远前者后面几个stage在做token之间的自注意力后者还在做局部卷积两个logits之间的软标签差异并不完全等价于“教师知道得更多”。RKD换了一个视角不要求学生复现教师的输出而是复现教师特征空间里样本间的几何关系。原论文用两个损失来描述这种关系。distance-wise loss先把样本间的负欧氏距离经过softmax转成一个“距离概率分布”蒸馏时让学生自己的距离概率分布去接近教师angle-wise loss则取三个样本组成一个夹角对cosine值做softmax后再对齐。两者都使用交叉熵衡量分布差距。实际操作时这两个损失都只作用在batch内的样本对上不需要任何结构匹配。这个区别带来了一个直接好处RKD不要求师生特征维数一致不要求结构对应。教师输出768维、学生输出512维甚至学生投影到任意维度都能计算样本间关系。后面我给的框架里师生统一走的都是768维embedding但这是为了方便对比并不是RKD本身的要求。相比之下FitNets那种直接对齐中间层的做法要处理通道匹配、空间位置匹配、stage选择一堆问题在异结构师生上非常折腾。实际训练中RKD与logits KD并不冲突。常见做法是CE loss保证任务正确性RKD loss保证几何结构如果数据集不大再叠一个0.5权重的logits KD三者不冲突。要注意的是RKD的目的是让学生“理解”教师的关系结构而不是替教师做特征对齐所以损失函数之间需要分好工。2.2 CoatNet当教师的选型理由自注意力与位置编码带来的关系信号CoatNetCoAtNet把卷积和自注意力揉在同一个网络里前几个stage用深度可分离卷积快速降分辨率后几个stage用TRATemporal Reduction Attention在降采样后的token序列上做全局交互。这种结构让它同时具备卷积的局部先验和自注意力的全局建模能力。位置编码被加在每个stage的token序列上使得特征除了“是什么”还带着“在哪个位置”的信息。这对蒸馏意味着什么教师模型的关系信号越丰富学生能从关系中学到的就越多。ResNet-50的stage4输出的2048维特征图已经足够抽象但没有显式位置编码也没有跨位置交互样本间距离主要由局部语义决定CoatNet的embedding不仅包含局部语义还包含自注意力重排后的全局关系样本间距离对“类别边界”更敏感。RKD的distance-wise损失偏向粗粒度特征关系——哪些样本聚在一起哪些样本应该远angle-wise损失偏向细粒度特征关系——三个样本在特征空间中的相对夹角。CoatNet恰好在这两个粒度上都比ResNet输出更可靠。从工程选型看用CoatNet当教师还有一个现实理由timm里已经可以直接加载ImageNet预训练权重构造教师模型只需十来行代码。它比再训练一个大ResNet当教师便宜也比Swin Transformer这种纯ViT当教师的收敛曲线更友好。实际跑下来CoatNet_0的全局特征和ResNet学生特征之间的RKD损失下降得比Swin-T要平稳可能因为两者都保留了卷积的归纳偏置关系矩阵的尺度差异没有大到让训练大幅震荡。选择教师时重点不是堆参数量而是看教师和学生的“结构差”能不能被关系蒸馏有效传递。可以参考下面这个思路教师选择结构差异RKD的预期收益更大的ResNet结构几乎一致学到的主要是容量差异关系结构趋同收益有限CoatNetCNN自注意力混合关系信号密度高RKD能跨结构传递收益最明显Swin Transformer纯自注意力、token交互重关系更强但分布尺度差异大需更小学习率对大batch更敏感2.3 RKD的适用边界什么时候不该硬上RKD关系蒸馏不是银弹我先划三个边界。第一数据量太小不要硬上。RKD监督的是batch内样本对关系一两万张图的训练集里每个类的样本方差太大关系矩阵本身就不稳定。此时优先用logits KD加一些特征对齐而不是让模型学一个带噪声的几何结构。第二单卡batch size低于32时RKD的效果会明显缩水原因在关联损失计算的是batch内关系分布样本少了关系矩阵稀疏可学习的信号太弱。第三学生模型如果连基础任务都还没学好RKD叠加上去只会拖慢收敛。我一般先把学生单独用CE训练到一定baseline再打开RKD做精修或者用warmup把RKD的权重从0拉到目标值。如果项目指标很紧、训练资源有限还有一个折中方案先用CoatNet的logits做一轮HKD让学生快速进入正确分类区间然后加载这版权重再启RKD精修。这个两步走方案比直接混训更容易定位问题也更容易在中期报告里解释收益来源。3. 蒸馏框架设计把ResNet改造成CoatNet的学生网络3.1 学生网络改造去掉分类头引出embedding既然RKD作用在特征空间就不能让ResNet只在末尾输出1000类logits。我一般把resnet50的最后一层全连接拿掉保留conv1到layer4的卷积骨架再接一个全局平均池化和一个线性投影头。投影头输出768维和CoatNet教师保持同一维度方便后续计算距离矩阵和角度矩阵。import torch import torch.nn as nn from torchvision.models import resnet50, ResNet50_Weights EMBED_DIM 768 class ResNetStudent(nn.Module): 改造后的ResNet50返回embedding和分类logits。 def __init__(self, num_classes100): super().__init__() backbone resnet50(weightsResNet50_Weights.IMAGENET1K_V2) self.features nn.Sequential( backbone.conv1, backbone.bn1, backbone.relu, backbone.maxpool, backbone.layer1, backbone.layer2, backbone.layer3, backbone.layer4, ) self.pool nn.AdaptiveAvgPool2d((1, 1)) # 用可训练的投影头把2048维压到768维 self.embed_head nn.Linear(2048, EMBED_DIM) # 分类头从embedding接出去保证任务loss和蒸馏loss共用同一套特征 self.classifier nn.Linear(EMBED_DIM, num_classes) def forward(self, x): h self.features(x) h self.pool(h).flatten(1) emb self.embed_head(h) logits self.classifier(emb) return emb, logits这里有个关键决定分类头接在embedding后面而不是直接用原来ResNet的1000维线性层。原因是RKD损失和CE损失必须作用在同一个特征空间如果分类头从原始2048维特征出学生就会倾向于把判别信息留在那2048维里embedding反而变成“应付蒸馏的旁路”。把两个头都接在768维embedding上可以让两个损失互相增益。AdaptiveAvgPool到1x1会丢掉空间位置信息。对RKD而言这不是问题因为关系计算发生在batch维度上不依赖空间坐标。如果你想保留更多空间细节可以改成AdaptiveAvgPool((7,7))再展平但线性层的参数会从49×2048变成约10万维训练初期投影头会很不稳定。先跑通最小可行版本再尝试空间信息增强是更务实的路径。3.2 教师网络加载CoatNet与timm的两种输出姿态教师这边直接用timm加载CoatNet预训练权重去掉分类头。需要特别留意timm不同版本对num_classes0的处理有差异有的版本直接返回全局池化后的向量[B,D]有的版本返回token序列[B,L,D]。CoatNet的forward_features在多数版本里返回的是最后一个TRA stage的token序列所以代码里要做一次维度判断避免被版本卡死。import timm class CoatNetTeacher(nn.Module): 封装CoatNet教师只取全局embedding。 def __init__(self, model_namecoatnet_0_rw_224): super().__init__() self.model timm.create_model(model_name, pretrainedTrue, num_classes0) torch.no_grad() def forward(self, x): tokens self.model.forward_features(x) # 兼容不同timm版本有的返回[B,L,D]有的已经池化到[B,D] if tokens.ndim 3: emb tokens.mean(dim1) else: emb tokens return emb使用torch.no_grad()是为了确保教师不会参与梯度回传。很多蒸馏代码跑着跑着显存爆炸就是因为漏了这一层教师参数在反向传播中被额外计算了一次更新GPU开销直接翻倍。这点在第5章还会展开。CoatNet的token序列经过mean后得到一个全局embedding。这里不需要特意选择TRA的某层输出用最后一级的token均值就够了。如果想更精细可以取出各层token做加权组合但那是进阶调优不建议第一个版本就引入这么多变量。3.3 组合成最小蒸馏单元两个模型定义好之后还需要把数据输入、优化器和损失函数串起来。我习惯用CIFAR-100做第一个蒸馏版本因为迭代快、显存压力小等RKD在中小数据集上符号验证跑通再平滑迁移到ImageNet。CIFAR-100原始分辨率只有32x32而CoatNet预训练权重是按224x224训练的所以需要先Resize到232再随机裁剪到224让教师看到的空间尺度不至于和预训练阶段差太多。from torchvision import transforms transform_train transforms.Compose([ transforms.Resize(232), transforms.RandomCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ])训练时有一个注意点CIFAR-100是从32x32上采样到224的上采样会引入插值模糊。这个问题在实际中比想象中严重因为CoatNet教师在ImageNet上从没见过这种模糊纹理教师给出的关系判断可能偏保守。我通常会建议把teacher也做一次简单的fine-tune或者至少用augmix做输入扰动让教师适应低分辨率上采样的域。如果不想额外训练教师那就在评估指标里给蒸馏基线留点余量别期待蒸馏后直接超过ImageNet精度。学生模型加载的是ImageNet预训练权重分类头输出改成100类。第一轮迭代时会发现CE loss下降比较快因为卷积骨架已经学会了通用特征这正好给了RKD一个良好的起点。整个最小蒸馏单元到这里已经齐了下一步就是把RKD损失函数写出来并接进训练循环。4. 训练循环与RKD损失核心代码和参数说明4.1 distance-wise损失负距离分布对齐RKD的distance-wise损失实现并不复杂。先把教师特征的样本间欧氏距离矩阵算出来乘以负号再除以温度然后按行做softmax得到一个“离我越近概率越大”的分布。学生特征也做同样操作最后用交叉熵让两个分布接近。下面这段代码直接可以跑import torch.nn.functional as F def rkd_distance_loss(s_feat, t_feat, temperature0.07): distance-wise RKD对齐样本间距离分布。 with torch.no_grad(): t_dist torch.cdist(t_feat, t_feat, p2) t_prob F.softmax(-t_dist / temperature, dim1) s_dist torch.cdist(s_feat, s_feat, p2) s_prob F.softmax(-s_dist / temperature, dim1) return torch.sum(-t_prob * torch.log(s_prob 1e-8), dim1).mean()torch.cdist计算的是两两欧氏距离形状为[B,B]。温度temperature控制分布的锐度温度越低距离近的样本对在分布中占的权重越大温度越高分布越均匀梯度信号越弱。原论文里distance-wise的默认温度是0.07我在CIFAR-100上一般从0.07开始如果发现loss下降太慢会试探性调到0.04左右但不会低于0.03否则分布会退化到几乎只看最近邻关系信息反而丢失。这里有个容易被忽略的细节距离矩阵对角线是0softmax之后对角线会拿到一个不小的概率相当于每个样本和自身的关系也在监督范围内。RKD原实现没有排除自身实践中保留它也没问题因为教师和学生对角线都是同样规则不会引入错误信号。如果你用的大batch有重复样本这个问题会更明显后续可以自己实现mask。4.2 angle-wise损失三元组夹角的分布对齐angle-wise损失计算的是三样本夹角。对每个锚点i样本j和样本k在特征空间里与i构成一个夹角取cosine值。cosine越大夹角越小三个样本在几何上越“挤在一起”。对每个锚点i所有(j,k)组合形成一个夹角矩阵再softmax成概率分布。def rkd_angle_loss(s_feat, t_feat, temperature0.08): angle-wise RKD对齐三元组夹角分布。 def angle_matrix(feat): # diff[i,j] feat[j] - feat[i] diff feat.unsqueeze(1) - feat.unsqueeze(0) # [B,B,D] norm diff.norm(dim-1, keepdimTrue).clamp_min(1e-12) unit diff / norm # 对每个锚点iunit[i]是一组方向向量 # 计算两两方向向量的点积得到cosine cos torch.einsum(ijk,ilk-ijl, unit, unit) # [B,B,B] return cos with torch.no_grad(): t_angle angle_matrix(t_feat) t_prob F.softmax(t_angle / temperature, dim-1) s_angle angle_matrix(s_feat) s_prob F.softmax(s_angle / temperature, dim-1) return torch.sum(-t_prob * torch.log(s_prob 1e-8), dim-1).mean()这段代码的计算量是O(B^3)因为对每个锚点都要生成一个BxB的夹角矩阵。B64时一次前向要算262144个夹角GPU能扛住B128时直接到200万量级训练速度会明显下降。我自己的经验是B64是性价比最高的点。如果显存不够优先降输入分辨率而不是降batch因为batch直接决定了关系分布的丰富程度。einsum那句是核心。unit形状是[B,B,D]对固定锚点iunit[i]是B个从i指向其他样本的单位向量unit[i]与unit[i]做点积得到夹角cosine矩阵。用ijk,ilk-ijl的写法j和l分别代表从i出发的两个方向最后一维做内积后落到i、j、l三个索引上正好形成[B,B,B]的夹角矩阵。这个写法比用torch.bmm直观得多不容易把维度弄反。4.3 训练循环与三个超参数训练循环的主体是常规的CE加两个RKD损失。教师全程eval学生trainAMP混合精度做前向和反向。需要注意student的forward会返回emb和logits两个值CE只用logitsRKD只用emb两者互不干扰。EPOCHS 60 alpha 1.0 temp_d 0.07 temp_a 0.08 for epoch in range(EPOCHS): student.train() teacher.eval() for images, targets in train_loader: images images.cuda(non_blockingTrue) targets targets.cuda(non_blockingTrue) optimizer.zero_grad() with torch.amp.autocast(cuda): s_emb, s_logits student(images) t_emb teacher(images) loss_ce F.cross_entropy(s_logits, targets) loss_dist rkd_distance_loss(s_emb, t_emb, temperaturetemp_d) loss_angle rkd_angle_loss(s_emb, t_emb, temperaturetemp_a) loss loss_ce alpha * (loss_dist loss_angle) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()使用torch.amp.autocast的前提是PyTorch 2.x如果还在用旧版本要换成torch.cuda.amp.autocast。CIFAR-100上60个epoch足够了ImageNet这种规模需要把epoch拉到100以上并且配合cosine schedule。三个超参数是alpha、temp_d、temp_a。它们的作用对象不同需要分开调参数我的常用值调节方向alpha1.0学生CE未收敛时从0 warmup到1.0防止RKD喧宾夺主temp_d0.07RKD loss不降时降到0.04训练震荡时回到0.1temp_a0.08同上但不要低于0.05否则角度分布过尖batch_size64小于32时关系采样不充分优先保证batchalpha的warmup特别重要。我一般前5个epoch设alpha0让学生先把CE任务学起来然后每epoch增加0.2到第10个epoch稳定在1.0。如果一上来就开满RKD学生要同时拟合任务边界和关系结构早期特征还没成形关系矩阵本身是乱的损失会一直卡在高位。5. 实战避坑与常见问题排查五个拿训练曲线换来的教训5.1 教师没断梯度显存直接翻倍现象接上RKD损失之后原本单卡16G能跑的batch直接OOM或者显存占用从9G跳到18G。原因新手写蒸馏代码最常见的错误是在计算RKD损失时把教师特征直接传进损失函数并且忘记detach。教师的forward虽然是前向计算但只要它输出参与了loss的backwardPyTorch就会为教师参数计算梯度并保留中间激活。CoatNet的token序列中间激活量比ResNet大得多OOM是必然结果。解决把教师forward整体包在torch.no_grad()里或者在损失函数内部用with torch.no_grad()锁定教师分支。我给出的CoatNetTeacher类已经处理了这个问题但如果你是自己魔改的教师封装记得在训练循环里断言一下t_emb.requires_grad为False用一行assert省得后面反复检查。5.2 RKD损失不降或反弹先查温度和权重现象distance-wise loss顺着epoch往下走但angle-wise loss一直横在2.3上下或者两个RKD loss在训练中期突然反弹变高。原因2.3正好是100类softmax均匀分布附近的交叉熵值说明angle矩阵的softmax分布接近均匀学生根本没从教师那边拿到有效梯度。这通常是两个原因叠加温度太高导致softmax过度平滑以及alpha权重太大让CE被迫让位学生特征退化到只有分类边界、没有几何结构。解决先把alpha降回0只看CE能不能正常下降。如果CE正常再单独调温度。对angle-wise损失我会把温度从0.08往0.05做网格搜索每次调整间隔至少10个epoch别在训练的前20个epoch里频繁换参数。同时把两个RKD损失分别记录到TensorBoard它们的下降节奏本来就不一样合在一起看容易误判。5.3 位置编码带来的对齐幻觉别拿MSE硬碰CoatNet现象试图在蒸馏框架里加一个中间层MSE损失把CoatNet某个TRA stage的token直接对齐ResNet某层feature map训练出来MSE很低但最终分类精度比不加还差。原因CoatNet的每个stage都带位置编码和自注意力交互它的token序列本身携带“位置语义”ResNet的feature map没有对应位置概念两者的空间坐标和语义坐标完全不是一回事。强迫学生去拟合带位置编码的token等于让学生学习一个它结构上表达不了的新任务。这种时候loss低不代表对齐得好反而说明学生把特征大量浪费在了迁就位置信息上。解决放弃逐位置对齐改用RKD的关系损失绕开位置对应。如果一定要加特征对齐只对class token或全局embedding做MSE并且权重控制在RKD的0.1倍以内。另一个思路是把ResNet的stage3和stage4输出类比FPN的多尺度特征各自投影到相同维度后分别计算距离矩阵再取加权平均作为RKD输入这样学生的多尺度信息也能被利用起来但这是进阶玩法第一个版本不建议碰。5.4 教师模型状态混用蒸馏特征漂移现象训练过程中RKD loss出现周期性波动每个epoch开头和结尾差一大截验证时RKD loss比训练时高一倍。原因CoatNet是从timm加载的预训练模型内部有DropPath和BatchNorm。如果教师模型误用了train()模式DropPath会随机丢弃路径BN也会继续统计running stats导致每次forward得到的教师特征粒度都不一样。教师本身是“正确答案”答案每天在变学生的蒸馏目标自然东倒西歪。解决训练循环的每个epoch开头强制执行teacher.eval()并且让student.train()继续保持。eval模式下DropPath关闭BN用固定的running stats教师输出才稳定。还有一个小细节如果用了torch.no_grad()可别再顺手调teacher.train()那会让DropPath重新打开no_grad只是不管梯度不代表模型切到推理模式。5.5 batch太小关系分布失真现象把batch从64降到16后RKD loss数值明显升高训练曲线抖动幅度变大最终精度比直接用CE还低。原因RKD的关系信号全部来自batch内部样本对。batch16意味着每个锚点只有15个邻居距离分布退化成几个最近邻的one-hot角度分布更是稀疏到几乎没有统计意义。梯度累积解决不了这个问题因为累积只是把多个小batch的梯度相加每个batch内部的关系多样性并没有增加。解决单卡batch至少32优先64。如果显存不够把输入分辨率从224降到192或者换显存更大的卡减少中间特征保存也可以比如把教师forward放到单独的stream里执行但别用降低batch来迁就显存。数据并行时每个GPU上的batch依然要保证32以上RKD损失只在单卡内计算跨卡的样本关系无法融合这是很多多卡实现容易忽视的点。6. 蒸馏效果验证不只比精度还要看关系一致性6.1 精度对比蒸馏前后必须同条件评估蒸馏模型的精度验证容易踩两个坑一是忘记切换eval模式导致BN统计和DropPath在推理阶段仍然处于训练状态二是用不同分辨率的transform去评估基线模型和蒸馏模型。我通常固定一套224x224的CenterCrop流程把直接训练的ResNet、CoatNet教师、蒸馏后的ResNet三者放在同一个验证集上对比。model.eval() correct 0 with torch.no_grad(): for images, targets in val_loader: images images.cuda(non_blockingTrue) _, logits model(images) pred logits.argmax(dim1) correct (pred.cpu() targets).sum().item() print(fTop-1 Acc: {correct / len(val_loader.dataset):.4f})如果蒸馏后的ResNet精度超过直接用CE训练的ResNet说明蒸馏有效如果只是持平也不一定失败还要看下面的关系一致性指标。原因很直接RKD优化的核心是关系结构而精度只是这个结构在分类边界上的投影两个模型可能精度相同但内部特征几何完全不同。6.2 距离矩阵一致性一个能说明问题的相关系数我习惯额外计算一个指标教师和学生特征矩阵的距离矩阵相关系数。具体做法是在验证集上各跑一遍把每个样本的embedding存下来形成两个[N,D]矩阵分别计算batch间的两两距离矩阵再拉平后求Pearson相关系数。相关系数越高说明学生的样本间距离排序越接近教师RKD的作用越明显。# 假设已经缓存了验证集特征 # t_feats/teacher_feats: [N,D]s_feats/student_feats: [N,D] t_dist torch.cdist(t_feats, t_feats) s_dist torch.cdist(s_feats, s_feats) corr torch.corrcoef( torch.stack([t_dist.flatten(), s_dist.flatten()]) )[0, 1].item() print(fDistance matrix correlation: {corr:.4f})这个指标在项目汇报里比top-1更有说服力因为它直接展示“学生学到了教师看数据的方式”。我通常会把不用RKD的ResNet作为对照组一起算如果加RKD后相关系数从0.65涨到0.85以上说明关系结构确实迁移过来了。最后分享一个训练习惯我会把验证集上的RKD损失单独记一条曲线作为早停指标。当val上的RKD loss还在下降而val acc不涨时说明蒸馏仍在改善学生与教师的关系一致性这时不急着停可以拉长训练当RKD loss开始回升时哪怕acc还在涨也该停了因为这往往意味着学生开始过拟合任务把已经学到的关系结构又冲淡了。这个习惯帮我避免过好几次“精度没涨但结构崩了”的翻车希望帮到你。本文还有配套的精品资源点击获取
