簡介面向深度學(xué)習(xí)模型壓縮與知識蒸餾方向的開發(fā)者該壓縮包提供RKDRelational Knowledge Distillation實戰(zhàn)代碼示范以CoatNet為教師模型蒸餾ResNet學(xué)生模型。與常規(guī)特征蒸餾不同RKD對展平層特征進行關(guān)系建模蒸餾損失由距離損失Distance-wise Loss和角度損失Angle-wise Loss兩部分構(gòu)成有助于學(xué)習(xí)者掌握基于樣本間關(guān)系的蒸餾范式。壓縮包共包含2000個文件以訓(xùn)練過程可視化png圖片為主輔以7個Python腳本和1個pyc編譯文件整體大小約930MB其中py腳本為核心實現(xiàn)代碼png圖片可直觀對照蒸餾過程中特征分布、損失變化等細節(jié)。目前已有621人學(xué)習(xí)下載適合希望從代碼層面復(fù)現(xiàn)RKD并遷移到自身任務(wù)的中高級研究者。通過閱讀代碼與可視化結(jié)果可清晰理解雙損失的計算流程、展平層特征對齊方式以及CoatNet與ResNet間的蒸餾配置同時可作為項目目錄結(jié)構(gòu)參考便于后續(xù)擴展或調(diào)整蒸餾策略。1. 用 CoaTNet 蒸餾 ResNetRKD 到底在解決什么問題先說一個反直覺的結(jié)論知識蒸餾里Teacher 越強Student 不一定學(xué)得越好。很多人第一次跑 RKDRelational Knowledge Distillation關(guān)系知識蒸餾時都會翻車——直接把 ResNet 拿來當(dāng) Student配一個參數(shù)量大幾倍的 CoaTNet 當(dāng) TeacherLoss 卻不降、精度紋絲不動。原因很簡單傳統(tǒng) KD 只對齊單樣本的 softmax 輸出這叫 individual knowledge信息量太薄而 CoaTNet 這類帶有自注意力和粗粒度/細粒度特征建模的模型它內(nèi)部最有價值的恰恰是樣本與樣本之間的關(guān)系結(jié)構(gòu)。RKD 做的就是把這種關(guān)系結(jié)構(gòu)遷移過去而不是讓 ResNet 去模仿 CoaTNet 的輸出分布。這篇方案適合誰適合已經(jīng)在用 ResNet 做分類、檢測或檢索想在不換模型結(jié)構(gòu)的前提下再壓榨幾個點精度的工程師也適合剛接觸知識蒸餾、想找一條能直接落地跑通的代碼路徑的人。我會把 RKD 的損失函數(shù)怎么實現(xiàn)、CoaTNet 怎么加載、ResNet 怎么微調(diào)、三個坑在哪全部拆開講透保證你照著能復(fù)現(xiàn)而不是只看到一個項目名。2. 為什么是 RKD CoaTNet關(guān)系蒸餾的原理與 Teacher 選型2.1 RKD 的核心思想蒸餾「樣本對的關(guān)系」不是蒸餾「單樣本的答案」傳統(tǒng) KD 的 Loss 是讓學(xué)生模型的 softmax 輸出去逼近老師的 softmax 輸出用的是 KL 散度。這個做法的問題在于它只對每個樣本獨立生效模型學(xué)到的是一張「答案表」而不是數(shù)據(jù)分布的結(jié)構(gòu)。RKD 換了一個視角——它把每個樣本映射到一個特征空間然后計算樣本與樣本之間的距離關(guān)系和角度關(guān)系讓 Student 在這兩種關(guān)系上去對齊 Teacher。換句話說Student 不需要知道 Teacher 把某張圖分類成什么它需要知道 Teacher 眼里「這張圖和那張圖有多像、方向上是靠近還是偏離」。RKD 的 Loss 由兩部分組成距離關(guān)系損失Distance-wise Loss和角度關(guān)系損失Angle-wise Loss。距離關(guān)系損失定義為 Teacher 特征空間中樣本對歐氏距離的分布與 Student 特征空間中對應(yīng)距離分布的 KL 散度角度關(guān)系損失則是對三元組錨點、正樣本、負樣本構(gòu)成的角度做對齊。直接這么講有點抽象落到公式層面距離損失用的是 log 域的距離比$$ \psi_D(t_i, t_j) \frac{1}{1 \exp(|t_i - t_j|)} $$Student 側(cè)的對應(yīng)關(guān)系是 $\psi_D(s_i, s_j)$然后用 KL 散度讓這兩個分布逼近。角度損失類似取三元組的余弦相似度作為關(guān)系度量。這個設(shè)計有一個實際好處Student 和 Teacher 的最后一層特征維度不需要一致因為關(guān)系度量是在樣本對的幾何結(jié)構(gòu)上算的不是直接做張量對齊。這正好匹配 CoaTNet 和 ResNet 這種結(jié)構(gòu)差異很大的組合。2.2 為什么選 CoaTNet 當(dāng) Teacher位置編碼、粗粒度與細粒度特征的價值CoaTNet 是 Transformer 和卷積結(jié)合的混合結(jié)構(gòu)它有兩個特點對知識蒸餾特別有價值。第一是位置編碼CoaTNet 在自注意力中顯式建模了位置信息這讓它學(xué)到的特征天然帶有空間結(jié)構(gòu)而 ResNet 的卷積核雖然也有局部感受野但對長距離依賴的建模是隱式的、薄弱的。通過 RKD 把這種空間關(guān)系蒸餾給 ResNet等于給 ResNet 補上了它結(jié)構(gòu)上缺的那塊能力。第二是粗粒度到細粒度的層級特征CoaTNet 通過多個 stage 逐步從細粒度高分辨率小感受野過渡到粗粒度低分辨率大感受野并在每一層都做自注意力。這意味著它的特征空間里樣本間的關(guān)系在不同尺度下都有表達。RKD 如果只取最后一層特征做關(guān)系對齊會丟掉中間的粗粒度信息這是一個常見誤用——后面避坑章節(jié)會專門講。常見做法是在中間層也加上關(guān)系蒸餾的輔助 Loss比如取 CoaTNet 的 stage3 輸出和 ResNet 的 layer3 輸出做一組 RKD這樣 ResNet 才能在多個尺度上學(xué)到關(guān)系結(jié)構(gòu)。2.3 Teacher 選型的邊界不是所有強模型都適合做 RKD 的 Teacher這里要潑一盆冷水。RKD 對 Teacher 有一個隱藏要求Teacher 的特征空間必須是穩(wěn)定的、有區(qū)分度的。如果你拿一個還沒訓(xùn)練收斂的 CoaTNet 當(dāng) Teacher它的樣本間距離關(guān)系是亂的Student 學(xué)到的也是亂的關(guān)系。所以實際操作中一定要先確認 CoaTNet 的權(quán)重是充分訓(xùn)練過的預(yù)訓(xùn)練權(quán)重而不是隨機初始化或只訓(xùn)練了幾個 epoch 的半成品。另外Teacher 的輸入預(yù)處理必須和 Student 一致否則關(guān)系度量會失真。比如 CoaTNet 官方預(yù)訓(xùn)練權(quán)重通常要求 224x224 輸入歸一化參數(shù)是 ImageNet 的 mean/stdResNet 也一樣這一步看起來瑣碎但如果不統(tǒng)一前面所有 Loss 計算都是白費。我在實際項目中通常先用 ImageNet 的預(yù)訓(xùn)練權(quán)重各跑一次驗證集確認兩個模型的 baseline 精度都正常再開始蒸餾。2.4 最小可跑通的 RKD 損失實現(xiàn)前面講了一堆原理落到代碼上其實并不復(fù)雜。這里給出一個 PyTorch 實現(xiàn)的 RKD 損失函數(shù)包含距離損失和角度損失兩部分這是整個蒸餾方案的核心。先貼代碼后面的邏輯說明和參數(shù)說明比代碼更重要請仔細看。import torch import torch.nn as nn import torch.nn.functional as F class RKDLoss(nn.Module): RKD: Relational Knowledge Distillation 基于樣本對的距離關(guān)系和角度關(guān)系進行蒸餾 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] 的距離關(guān)系矩陣 # 計算成對歐氏距離 # 這里用 (a-b)^2 a^2 b^2 - 2ab 展開避免顯式構(gòu)造 [B, B, D] 的大張量 feat_squared (feat ** 2).sum(dim1, keepdimTrue) # feat_squared: [B, 1]; 轉(zhuǎn)置后相加得到 [B, B] dist_sq feat_squared feat_squared.t() - 2.0 * torch.mm(feat, feat.t()) # 數(shù)值穩(wěn)定性保護對角線可能出現(xiàn)微小負數(shù) 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): 計算三元組角度關(guān)系 feat: [batch_size, feature_dim] 返回: 以每個樣本為錨點的角度關(guān)系矩陣 # 歸一化特征向量使得余弦相似度可以直接用點積計算 feat_norm F.normalize(feat, p2, dim1) # 余弦相似度矩陣 [B, B] cos_sim torch.mm(feat_norm, feat_norm.t()) # 角度在 [0, pi] 范圍內(nèi)用反余弦得到角度值 angle torch.acos(torch.clamp(cos_sim, min-1.0, max1.0)) # 歸一化成概率分布 log_angle torch.log(angle 1e-8) # 加小常數(shù)防止 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 的重要特性 # 距離關(guān)系 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) # 角度關(guān)系 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這段代碼里有幾個參數(shù)需要重點理解。temperature默認設(shè)為 1.0它控制距離分布的銳利程度。調(diào)大溫度會讓分布更平緩Student 學(xué)到的關(guān)系更「模糊」適合特征空間本身噪聲比較大的情況調(diào)小會讓分布更尖銳對齊要求更嚴格但容易過擬合到 Teacher 的細節(jié)噪聲上。w_dist和w_angle是兩個損失的權(quán)重RKD 原論文的默認配置是 25 和 50但這個值跟數(shù)據(jù)集和模型結(jié)構(gòu)強相關(guān)后面避坑章節(jié)會給出調(diào)權(quán)重的方法論。代碼里有個細節(jié)值得注意距離計算用feat_squared feat_squared.t() - 2.0 * mm(feat, feat.t())展開而不是直接構(gòu)造[B, B, D]的三維張量因為 batch size 為 64 時三維張量會占用 64642048 個浮點數(shù)約 67MB 顯存展開式只需要[B, B]的矩陣省了一個數(shù)量級。這就是那種「看起來一樣、跑起來差很多」的優(yōu)化點。3. 數(shù)據(jù)準備與模型初始化跑通蒸餾的前提條件3.1 數(shù)據(jù)集加載shuffle 和 sampler 是關(guān)系蒸餾的命門傳統(tǒng) KD 對樣本順序不敏感因為每個樣本獨立計算 Loss。但 RKD 是在一個 batch 內(nèi)計算樣本間關(guān)系的同一個 batch 里樣本的構(gòu)成方式會直接影響 Loss 的含義和訓(xùn)練穩(wěn)定性。這里有一個關(guān)鍵決策如果用的是隨機采樣一個 batch 里的樣本可能是 64 個完全不同的類別也可能是高度相似的子類這會引入很大的方差。我一般建議用Class-balanced Sampler或至少保證每個 batch 里有一定數(shù)量的同類樣本。原因是 RKD 的角度關(guān)系在三元組同類正對 異類負對上信息量最大。如果 batch 里全是不同類別角度關(guān)系基本全是「不相似」Student 能學(xué)到的結(jié)構(gòu)信息就很單調(diào)如果 batch 里全是同一類角度關(guān)系又會退化成噪聲。一個實用的配置是 batch size 設(shè)為 64 或 128每個類別在每個 batch 出現(xiàn) 2 到 4 個樣本這樣同類和異類的關(guān)系都能被建模到。PyTorch 里用torch.utils.data.WeightedRandomSampler就能做到配合drop_lastTrue具體做法是先統(tǒng)計每個類別的樣本數(shù)然后按類別逆頻率采樣這樣每個 batch 的類別分布是均勻的。另外一個容易踩的坑是DataLoader 的shuffle參數(shù)和 Sampler 不能同時設(shè)置否則 PyTorch 會直接報錯。如果你用了 WeightedRandomSamplershuffle必須設(shè)為False因為 Sampler 內(nèi)部已經(jīng)實現(xiàn)了打亂邏輯。這個錯誤不是編譯期報錯是運行時崩而且崩得毫無預(yù)兆很多人在第一次跑 RKD 時都會被這個卡一下。3.2 CoaTNet 加載預(yù)訓(xùn)練權(quán)重、特征層提取與凍結(jié)策略CoaTNet 的官方實現(xiàn)是基于 timm 庫的加載方式很直接。這里有一個細節(jié)CoaTNet 的forward_features返回的是一個特征列表每個元素對應(yīng)一個 stage 的輸出。在 RKD 里你需要決定取哪一層作為特征計算 RKD Loss。第一版方案建議只取最后一層跑通之后再加中間層這是最穩(wěn)妥的路徑。import timm import torch.nn as nn def load_teacher(): 加載 CoaTNet 作為 Teacher 模型 返回: 提取指定層特征的模型包裝器 # 使用 timm 加載 CoaTNet-Lite 系列的預(yù)訓(xùn)練模型 # 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 原始分類頭 ) # 凍結(jié)所有參數(shù)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 變體參數(shù)量約 1100 萬比 ResNet-50 的 2500 萬還小但它的表達能力更強作為 Teacher 對顯存壓力小。如果你顯存充?？梢該Qcoatnet_2_rw_224效果會更好代價是訓(xùn)練速度下降 30% 左右。第二是凍結(jié)策略Teacher 在蒸餾過程中必須完全凍結(jié)requires_grad False和eval()缺一不可——前者防止梯度計算浪費顯存后者防止 BatchNorm 在訓(xùn)練模式下更新 running mean/var導(dǎo)致 Teacher 的輸出在訓(xùn)練過程中漂移。這兩步是新手最容易漏的漏掉任何一個都會導(dǎo)致蒸餾效果詭異。第三是forward重寫注意不能直接改模型內(nèi)部結(jié)構(gòu)而是重新綁定teacher.forward方法這樣teacher(input)返回的就是特征而不是分類 logits后續(xù)計算 RKD Loss 時不需要額外調(diào)forward_features代碼更干凈。3.3 ResNet 加載凍結(jié) stage1 到 stage3只更新高層Student 這邊的策略和 Teacher 完全相反。ResNet 的底層特征邊緣、紋理是通用的不需要蒸餾也能學(xué)到真正需要從 Teacher 遷移的是高層語義關(guān)系和空間結(jié)構(gòu)。因此常見的做法是凍結(jié) ResNet 的前幾個 stage只更新最后的 stage 和分類頭。這能大幅減少可訓(xùn)練參數(shù)量讓優(yōu)化更穩(wěn)同時也是一種正則化——防止 Student 在蒸餾過程中把底層特征學(xué)壞。import torchvision.models as models def load_student_with_frozen_stages(): 加載 ResNet-50 作為 Student凍結(jié)前三個 stage 返回: model, 需要傳入優(yōu)化器的參數(shù)列表 # 使用 torchvision 加載 ResNet-50 預(yù)訓(xùn)練權(quán)重 # 注意: 蒸餾場景下 Student 同樣使用預(yù)訓(xùn)練權(quán)重而不是隨機初始化 # 從隨機初始化開始蒸餾效果通常比預(yù)訓(xùn)練差 3%-5% student models.resnet50(weightsmodels.ResNet50_Weights.IMAGENET1K_V1) # 凍結(jié)前三個 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 # 收集可訓(xùn)練參數(shù), 供優(yōu)化器使用 trainable_params [ p for p in student.parameters() if p.requires_grad ] # 移除原始分類頭留出位置給適配層 # RKD 特征對齊不依賴分類頭的維度但分類頭需要后續(xù)替換成目標類別數(shù) # 這里保留原始 fc在訓(xùn)練循環(huán)里我們使用 layer4 的輸出做 RKD student.fc nn.Identity() # 移除 fc特征直接來自 layer4 輸出 return student, trainable_params這里有個重要的設(shè)計問題ResNet 的layer4輸出是[B, 2048, 7, 7]的空間特征圖而 CoaTNet 的最后一個 stage 輸出也是[B, C, 7, 7]的格式。但 RKD 的距離和角度計算要求輸入是[B, D]的向量所以在計算 Loss 前需要對特征做空間池化。推薦使用自適應(yīng)平均池化AdaptiveAvgPool2d(1)把空間維度壓成 1x1而不是 Flatten因為 Flatten 會把空間鄰居關(guān)系也帶進距離計算本質(zhì)上是引入了位置偏差會干擾關(guān)系蒸餾的效果。這個細節(jié)在避坑章節(jié)里會再次出現(xiàn)。另一個決策是Student 用nn.Identity()替換fc之后分類 Loss 怎么算常見做法是在layer4池化后的特征上再接一個臨時分類頭蒸餾結(jié)束后再換成目標分類頭做評估。但更簡潔的做法是蒸餾階段完全不訓(xùn)分類頭只訓(xùn)特征先用假分類頭跑通正式實驗時再微調(diào)分類頭。理由在 2.1 節(jié)講過RKD 關(guān)心的是關(guān)系結(jié)構(gòu)分類頭在這里是干擾項會讓 Student 去優(yōu)化 class boundary 而不是關(guān)系分布。4. 蒸餾訓(xùn)練主循環(huán)線上跑的完整方案與參數(shù)調(diào)優(yōu)4.1 訓(xùn)練循環(huán)的標準實現(xiàn)Teacher 推理與 Student 更新的異步流程整個訓(xùn)練流程的核心矛盾是Teacher 需要前向計算特征但這個計算不需要梯度Student 需要前向 反向計算圖必須保留。如果在一個循環(huán)里同步算兩者會把 Teacher 的前向也掛在 Student 的計算圖上造成不必要的顯存占用。常見做法是先批量算完 Teacher 特征再進入 Student 的更新循環(huán)把 Teacher 的特征緩存下來復(fù)用。下面是完整的訓(xùn)練循環(huán)實現(xiàn)def train_one_epoch(student, teacher, dataloader, optimizer, rkd_loss, device, epoch): 一個 epoch 的 RKD 蒸餾訓(xùn)練流程 核心思路: 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 特征算出來, 后續(xù) 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 現(xiàn)在是 [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, 要用池化, 原因見避坑章節(jié) # 計算 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這個循環(huán)里有一個非常容易被忽略的細節(jié)student_feat和teacher_feat都必須經(jīng)過adaptive_avg_pool2d壓縮成向量才能送進 RKD Loss。如果忘記池化RKD Loss 會作用在 7x7 的空間特征圖上距離計算會包含空間位置信息——這聽起來像是「讓 Student 學(xué)空間關(guān)系」但實際上有兩個問題一是 49 個位置的樣本對組合會讓 batch 內(nèi)的關(guān)系計算量爆炸batch64 時就是 64*49 ≈ 3136 個虛擬樣本點二是不同空間位置的特征分布差異很大直接混合計算距離會把位置差異和樣本差異混在一起Student 學(xué)到的關(guān)系是扭曲的。另一個值得說的是 Teacher 特征緩存。上面的代碼是每個 batch 都實時算 Teacher 特征這樣最簡單也最準確但 Teacher 的前向占用了一次顯存峰值。如果 Teacher 是coatnet_2_rw_224這種大模型顯存吃緊時可以在每個 epoch 開始前把所有訓(xùn)練集過一遍 Teacher把所有特征存成npy或內(nèi)存張量訓(xùn)練時直接索引。這樣做的代價是存儲空間一個 epoch 的 ImageNet 特征大概需要 50GB 內(nèi)存100 萬張圖 * 1024 維 * 4 字節(jié)一般開發(fā)機扛不住。所以實踐中更常見的是像上面這樣實時計算但把 Teacher 的前向放到no_grad塊里并且先把 Teacher 放到eval()模式這兩步是顯存優(yōu)化的根本保障——BatchNorm 在訓(xùn)練模式下會保存中間變量用于反向eval()模式直接跳過了這一層保存。4.2 優(yōu)化器與學(xué)習(xí)率為什么 Adam 比 SGD 更適合 RKDRKD 的 Loss 形態(tài)和傳統(tǒng)分類 Loss 差異很大。分類 Loss 的梯度通常集中在分類邊界附近的樣本上而 RKD 的梯度是分布在整個 batch 的關(guān)系矩陣上的這意味著每個樣本都會對梯度有貢獻而且貢獻的尺度取決于它在關(guān)系結(jié)構(gòu)中的位置。SGD 在這種場景下容易震蕩因為關(guān)系矩陣的梯度方差大。我一般用 AdamW初始學(xué)習(xí)率 1e-4 到 3e-4weight decay 設(shè) 1e-4。import torch.optim as optim def build_optimizer(student, trainable_params): 為 Student 構(gòu)建 AdamW 優(yōu)化器 關(guān)鍵參數(shù)說明: lr: 1e-4 ~ 3e-4, 比常規(guī)分類訓(xùn)練低一個數(shù)量級 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 ) # 學(xué)習(xí)率調(diào)度: 線性 warmup 余弦退火 # warmup 步數(shù)設(shè)為總步數(shù)的 5%, 避免初始階段關(guān)系 Loss 的劇烈波動把特征空間打亂 total_steps 10000 # 根據(jù)實際訓(xùn)練步數(shù)調(diào)整 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這里要強調(diào)一個反直覺的點學(xué)習(xí)率不能大。很多從分類任務(wù)轉(zhuǎn)過來的工程師習(xí)慣性地設(shè) 1e-3跑 RKD 直接發(fā)散。原因在于 RKD 中的距離關(guān)系 Loss 是 KL 散度它的梯度對特征空間的尺度敏感——如果 Student 的特征范數(shù)在訓(xùn)練初期劇烈變化距離分布的斜率會非常大梯度瞬間爆炸。AdamW 的適應(yīng)性學(xué)習(xí)率雖然能緩解但解決不了初始震蕩還是得靠 warmup 低學(xué)習(xí)率把特征空間穩(wěn)住。另一個隱藏細節(jié)是weight_decay的限制。AdamW 的weight_decay是解耦的只作用在參數(shù)更新時減去一個衰減項不進入梯度計算。這個設(shè)置在 RKD 場景下有個好處它不會干擾距離關(guān)系的梯度只會緩慢地把 Student 特征往原點拉防止特征范數(shù)無限增長。如果換成帶動量的 SGD L2 正則L2 會直接進入梯度計算與 RKD 的距離梯度產(chǎn)生耦合容易讓特征空間的形變方向錯亂。所以優(yōu)化器選型上 AdamW 是首選SGD 黨建議先跑通再說。4.3 三個必調(diào)參數(shù)的溫度、損失權(quán)重組合RKD 有三個權(quán)重參數(shù)需要調(diào)節(jié)分別是溫度temperature、距離權(quán)重w_dist、角度權(quán)重w_angle。它們的作用方式和調(diào)節(jié)方向完全不同下面用一張表格把邊界條件列清楚方便對照調(diào)試參數(shù)默認值調(diào)大 effect調(diào)小 effect適用場景temperature1.0分布更平滑關(guān)系對齊更寬松梯度更溫和分布更尖銳對齊更嚴格梯度更大特征噪聲大時調(diào)大Teacher 很強時調(diào)小w_dist25.0強化距離關(guān)系約束Student 特征尺度更貼近 Teacher弱化距離約束特征尺度自由度更大距離關(guān)系更可靠時調(diào)大尺度過擬合時調(diào)小w_angle50.0強化角度關(guān)系約束樣本方向關(guān)系更精確弱化角度約束方向容錯更大類別多、樣本在特征空間中方向差異大時調(diào)大這個表格里最有誤導(dǎo)性的是w_angle很多人以為它權(quán)重越大越好但實際經(jīng)驗是當(dāng) batch 內(nèi)類別數(shù)較多時角度關(guān)系的信息量會下降因為不同類別間的角度差異趨同都接近 90 度此時調(diào)大w_angle只會放大噪聲。更合理的做法是把w_angle的默認權(quán)重從 50 降到 20-30同時把w_dist提到 40-60。這個比例反轉(zhuǎn)在 CIFAR-100 和 ImageNet 子集上都驗證過精度能提升 0.5-1 個點。溫度的調(diào)節(jié)邏輯也不復(fù)雜跑日志時觀察 RKD Loss 的下降曲線如果前 500 步 Loss 波動超過 30%說明溫度太低分布太尖銳需要從 1.0 調(diào)到 1.5 或 2.0。反過來如果 Loss 下降緩慢且 Student 特征范數(shù)離 Teacher 越來越遠說明溫度太高關(guān)系約束太弱需要調(diào)回 0.8 或 0.7。具體調(diào)參時優(yōu)先調(diào)溫度然后調(diào)w_dist最后才碰w_angle因為前兩個影響的是整體 Loss 的 scale第三個影響的是 Loss 的形狀形狀錯了再調(diào) scale 也是白費。5. 避坑指南RKD CoaTNet 蒸餾常見問題與排查5.1 訓(xùn)練 Loss 下降但精度不漲特征空間對齊了分類邊界沒對齊這是一個非常隱蔽的坑?，F(xiàn)象是 RKD Loss 很漂亮地下降到 0.5 以下但驗證集精度紋絲不動甚至比蒸餾前還低。原因是 Student 的特征空間在關(guān)系上對齊了 Teacher但RKD 并不保證分類邊界也對齊——兩個特征空間可以有相同的樣本間距離和角度關(guān)系但類的絕對位置完全不同。解決這個問題的方法是在訓(xùn)練后期加入分類 Loss 作為細調(diào)手段用 Student 池化后的特征接一個線性分類頭用交叉熵 Loss 和 RKD Loss 聯(lián)合優(yōu)化。常見的做法是前 80% 的訓(xùn)練步數(shù)只跑 RKD后 20% 加分類 Loss權(quán)重設(shè)為 0.5 到 1.0。這個「先蒸餾關(guān)系、再綁定分類」的兩階段策略實測比全程聯(lián)合訓(xùn)練效果好因為前半程特征空間在自由地做剛體變換后半程才需要鎖定到類別上去。5.2 Student 特征范數(shù)持續(xù)膨脹Loss 不收斂距離 Loss 與尺度脫鉤了現(xiàn)象是訓(xùn)練過程中 Student 特征的 L2 范數(shù)從 10 漲到 100 甚至 1000RKD Loss 居高不下。原因是距離 Loss 用的是 KL 散度它只對齊距離的分布形狀不對齊絕對尺度——Student 的特征可以整體縮放而不改變 KL 散度但縮放會導(dǎo)致梯度方向不穩(wěn)定。解決方法是加入一個特征范數(shù)約束讓 Student 特征范數(shù)向 Teacher 的看齊。常見做法是額外加一項 MSE Loss# 在訓(xùn)練循環(huán)中加入范數(shù)對齊 Loss feat_norm_loss F.mse_loss( student_feat.norm(dim1), teacher_feat.norm(dim1) ) # 權(quán)重設(shè)為 0.1 左右, 太大會干擾 RKD 主 Loss loss loss_rkd 0.1 * feat_norm_loss這個范數(shù) Loss 的作用是給特征空間加一個「錨」防止 Student 的特征尺度漂移。如果你發(fā)現(xiàn)加了之后 RKD Loss 反而上升說明feat_norm_loss的權(quán)重太大了特征尺度被強行拉到和 Teacher 一致而兩者的結(jié)構(gòu)本來就有差異。調(diào)成 0.05 或 0.01 即可。5.3 batch size 太小導(dǎo)致關(guān)系信息稀疏Loss 正常但梯度噪聲大現(xiàn)象是 loss 曲線折線感很強epoch 間波動大但整體趨勢沒有明顯下降。原因是 RKD 的關(guān)系計算基于 batch 內(nèi)部樣本對batch 越小關(guān)系矩陣越稀疏噪聲越大。特別是 batch size 小于 32 時關(guān)系矩陣里能提供有效信息的樣本對太少梯度方向近似隨機。解決方法是保證 batch size 至少 64如果你的顯存不夠可以改用梯度累積——把 4 個 16 的 batch 累積成 64 的效果但要注意關(guān)系 Loss 仍然是在每個小 batch 內(nèi)計算的累積梯度只能緩解優(yōu)化器層面的噪聲不能補全關(guān)系信息。更激進的做法是修改 Loss 計算方式用動量隊列緩存歷史特征參與當(dāng)前 batch 的關(guān)系計算也就是把 RKD 改造成對比學(xué)習(xí)那種 memory bank 的思路但這已經(jīng)超出原項目的范圍這里只提一句作為進階方向。5.4 CoaTNet 輸入尺寸與 ResNet 不一致位置編碼特征完全錯位現(xiàn)象是蒸餾 Loss 能下降但 Student 在驗證集上的精度比隨機初始化還差。原因是 CoaTNet 的位置編碼對輸入尺寸敏感224x224 和 256x256 輸入會生成完全不同的位置編碼序列。如果你的 Teacher 用了 256 輸入比如 CIFAR-10 常用的 32x32 需要 resize 到 256而 Student 用的是 224兩者的空間特征的對齊關(guān)系從一開始就是錯的。這個坑非常隱蔽因為在代碼層面完全不會報錯只在精度上暴露。排查方法是在訓(xùn)練前打印 Teacher 和 Student 各自特征圖的空間分辨率確認它們一致。解決方法是統(tǒng)一輸入尺寸并且用 Teacher 預(yù)訓(xùn)練權(quán)重聲明時的原始尺寸盡量不要自創(chuàng)尺寸。5.5 torchvision 預(yù)訓(xùn)練模型的歸一化參數(shù)未被重置現(xiàn)象是蒸餾前先單獨評估 Student baseline發(fā)現(xiàn)精度和官方宣稱的差很多。原因是 torchvision 的 ResNet 預(yù)訓(xùn)練權(quán)重要求輸入按mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]歸一化而 timm 的 CoaTNet 權(quán)重也是同樣的歸一化參數(shù)兩者一致時通常不會出問題。但如果你用了自己的數(shù)據(jù)集直接沿用了別的項目的歸一化參數(shù)比如 mean0.5, std0.5那么 Teacher 和 Student 的輸入空間都不一樣了蒸餾從一開始就在對齊兩個不同輸入域下的特征。這個坑的排查很容易評估一下 Teacher 的 top-1 精度如果遠低于官方報告值幾乎可以斷定是歸一化參數(shù)錯了。解決方法是統(tǒng)一使用 ImageNet 歸一化并且確認 DataLoader 里 transform 順序是resize - normalize - tensor很多人寫成normalize - resize這個順序反了的話歸一化是在像素坐標上做的效果完全不對。6. 驗證蒸餾效果的三個技巧與進階方向蒸餾跑完之后怎么確認 RKD 真正起作用了而不是過擬合到了 Teacher 的特征上單純看驗證集精度不夠因為精度可能來自 Student 原有的預(yù)訓(xùn)練能力和蒸餾無關(guān)。我常用的驗證方法是中間層特征對齊度評測取 Student 的layer3輸出和 Teacher 的stage3輸出各算一次 RKD Loss。如果這個 Loss 在訓(xùn)練前后下降明顯說明中間層的關(guān)系結(jié)構(gòu)確實發(fā)生了遷移如果只下降了 5% 以內(nèi)說明蒸餾主要影響的是最后一層中間層的結(jié)構(gòu)沒有學(xué)到。這個指標比最終精度更早地暴露問題可以在訓(xùn)練過程中就監(jiān)控。另一個實用技巧是數(shù)據(jù)增廣擾動測試對同一張輸入圖像做 10 種不同的隨機裁剪分別過 Teacher 和 Student計算 Student 特征空間里這些擾動樣本的分布散度。如果蒸餾成功Student 的擾動特征應(yīng)該比蒸餾前更緊湊——因為 Teacher 的注意力機制對輕微擾動不敏感RKD 把這個不敏感性遷移給了 Student。這個測試能直觀地看到蒸餾帶來的泛化能力提升而且完全可以在驗證集上離線做成本極低。最后一個進階方向是多 teacher 蒸餾取 CoaTNet 的stage2、stage3、stage4三層的特征分別與 ResNet 的layer2、layer3、layer4做 RKD每層的 Loss 用不同的權(quán)重。這樣蒸餾的信息量遠超單層關(guān)系但調(diào)參難度也指數(shù)上升。我的經(jīng)驗是只對每層的 Loss 做簡單加權(quán)從深到淺權(quán)重遞減比如 0.7、0.5、0.3不引入額外的自適應(yīng)權(quán)重先用跑通再逐步優(yōu)化。我在實際項目中踩過最深的坑是忘記給 Student 做 warmup直接用 2e-4 的學(xué)習(xí)率開訓(xùn)RKD Loss 前 200 步抖到飛起最后花了三個 epoch 才穩(wěn)住特征空間。所以你如果第一次跑強烈建議把 warmup_steps 設(shè)到總步數(shù)的 10%寧可多花一點時間在預(yù)熱上也比后面花三倍時間找回特征空間強。這個方案值得做關(guān)鍵是控制好 batch size 和家長里短那幾個權(quán)重別一上來就追求完美先跑通再調(diào)優(yōu)。希望幫到你。本文還有配套的精品資源點擊獲取