Markdownで表示

知識蒸留

蒸留は、生徒モデルの中間特徴マップを凍結済み教師モデルの特徴マップに近づける第2の損失項を追加します。LibreYOLOは順伝播フックで特徴量を取得するため、教師自身のヘッドと損失は関与しません。

大規模チェックポイントから蒸留

distill_modelを設定すると蒸留が有効になります。値には教師のチェックポイントを指定し、他のモデルと同じファクトリで読み込みます。

Python
from libreyolo import LibreYOLO # 同じファミリーの大きなチェックポイントで小さなモデルを指導するmodel = LibreYOLO("LibreYOLO9s.pt")model.train(    data="my-dataset.yaml",    epochs=100,    distill_model="LibreYOLO9c.pt",    distill_loss_type="mgd",)
CLI
libreyolo train model=LibreYOLO9s.pt data=my-dataset.yaml \  epochs=100 distill_model=LibreYOLO9c.pt distill_loss_type=mgd

教師はno_grad内で順伝播し、AMPが有効な場合はautocast内でも実行されるため、凍結済みモデルが各ステップで完全精度の計算コストを負うことはありません。順伝播フックは名前付きの取得点で特徴マップを取得し、損失が生徒の特徴量と比較します。その結果は学習損失へ加算され、distillという名前の構成要素として報告されます。

凍結済み基盤バックボーンから蒸留

代わりに、自己教師ありViTで生徒の1つのバックボーンステージを指導できます。教師の特徴量はフックではなく教師自身の特徴抽出器から取得し、損失がパッチグリッドと畳み込みストライドの不一致を処理します。

Python
from libreyolo import LibreYOLO # 凍結した自己教師ありViTでバックボーンの1ステージを指導するmodel = LibreYOLO("LibreYOLO9s.pt")model.train(    data="my-dataset.yaml",    epochs=100,    distill_model="dinov2",)
CLI
libreyolo train model=LibreYOLO9s.pt data=my-dataset.yaml \  epochs=100 distill_model=dinov2

distill_modelはDINOv2-baseを表すdinov2に加え、dinov2_vits14dinov2_vitb14dinov2_vitl14dinov2-smalldinov2-basedinov2-large、およびfacebook/dinov2で始まる任意の生のHub IDを認識します。それ以外は教師チェックポイントのパスとして扱われます。

この経路ではdistill_loss_typeにかかわらずfeat_mseを使用し、transformersのインストールが必要です。重みキーが不足した状態で読み込まれた教師は、部分的にランダムなバックボーンで蒸留せず処理を中止します。

対応ファミリー

蒸留への対応は生徒モデルのメソッドとして実装され、2種類あります。

get_distill_config()は検出器の教師が指導するマルチスケールの取得点を提供します。YOLOv9、YOLOX、RF-DETRが実装しています。

get_backbone_distill_config()は基盤モデルの教師が指導する1つのバックボーンステージを提供します。実装しているのはYOLOv9だけです。

それ以外では、損失なしで学習せず例外を送出します。

LibreDFINE does not implement get_distill_config(). Distillation is not yet
supported for the 'dfine' family.
Foundation-model distillation into the 'yolox' family is not supported yet
(no get_backbone_distill_config()).

特徴量の取得点

取得点はファミリーと役割ごとに固定されます。教師と生徒が同じアーキテクチャである必要はありませんが、特徴量のストライドは一致する必要があります。

ファミリー役割取得点ストライド
YOLOv9教師または生徒neck.elan_up2neck.elan_down1neck.elan_down28、16、32
YOLOv9基盤モデルの生徒backbone.elan316
YOLOX教師または生徒backbone.C3_p3backbone.C3_n3backbone.C3_n48、16、32
RF-DETR教師または生徒model.backbone.0.projector.stages.0セットアップ時に調査

ストライドが一致しない場合は、学習開始前に例外を送出します。

Teacher and student must have matching strides. Teacher: [8, 16, 32],
Student: [16]

基盤モデルの教師では、この確認を行いません。グリッドが異なること自体がこの経路の目的だからです。

3種類の損失

検出器の教師では、distill_loss_typeで特徴量の損失を選択します。基盤モデルの教師は常にfeat_mseを使用します。

mgd(Masked Generative Distillation)は、生徒の空間位置の一部をマスクし、小さな2層の畳み込みジェネレーターを学習して、残りの情報から教師の完全な特徴マップを再構築します。distill_mask_ratioはマスクする割合で、デフォルトは0.65です。

cwd(Channel-wise Distillation)は、各チャンネルの空間活性化を確率分布へ変換し、チャンネルごとにKLダイバージェンスを最小化します。distill_tauはsoftmax温度で、デフォルトは1.0です。

feat_mseは1×1畳み込みで生徒のチャンネルを教師に合わせ、教師のグリッドを生徒のグリッドへ双線形補間して、平均二乗誤差を取ります。distill_normalize=Trueは、最初に両方の特徴マップをチャンネル次元に沿ってL2正規化します。これにより比較は角度だけに基づき、スケールに依存しなくなります。デフォルトはFalseです。

disは全体に適用する重みです。未指定の場合、各損失が公開済みのデフォルト値を使用します。MGDは2e-5、CWDは1.0、特徴量MSEは1.0です。5桁も異なるため、ある損失用に調整した重みは別の損失では意味を持ちません。

損失を調整
from libreyolo import LibreYOLO model = LibreYOLO("LibreYOLO9s.pt")model.train(    data="my-dataset.yaml",    distill_model="LibreYOLO9c.pt",    distill_loss_type="cwd",    dis=1.0,           # 蒸留全体の重み    distill_tau=1.0,   # CWDのsoftmax温度)

distill_mask_ratiodistill_taudistill_normalizeにはCLIフラグがありません。Python引数またはcfg=のYAMLキーとして指定します。RF-DETRのCLI引数マッピングには蒸留キーがないため、RF-DETRの蒸留全体もPython専用です。

アダプター、チェックポイント、マルチGPU

各損失は生徒の外側に小さな学習可能モジュールを構築します。1×1チャンネルアダプターとMGDのジェネレーターです。これらには実行時の有効学習率を持つ専用のオプティマイザーパラメーターグループが割り当てられます。

これらのモジュールはチェックポイントのdistillerキーに書き込まれ、再開時に復元されます。そのため、再開した実行でプロジェクターが初期状態に戻ることはありません。

DDPではアダプターがラップ済みの生徒モデル外部にあるため、DDPリデューサーはその勾配を認識しません。トレーナーは各ステップで明示的にall-reduceを行い、すべてのランクで同じアダプターを学習します。

蒸留実行ではCUDAグラフのキャプチャを利用できません。cuda_graph=Trueを渡すとログを1行出し、通常実行で学習します。学習パフォーマンスを参照してください。

関連項目

LibreYOLO v1.5.0で検証済みです。