yoshitomo-matsubara/torchdistill

A coding-free framework built on PyTorch for reproducible deep learning studies. PyTorch Ecosystem. 🏆26 knowledge distillation methods presented at TPAMI, CVPR, ICLR, ECCV, NeurIPS, ICCV, AAAI, etc are implemented so far. 🎁 Trained models, training logs and configurations are available for ensuring the reproducibiliy and benchmark.

torchdistill – 以設定驅動的知識蒸餾框架

是什麼

一個基於 PyTorch 的開源 Python 庫，讓您無需撰寫自訂訓練迴圈即可執行 知識蒸餾 實驗（教師-學生模型訓練）。
所有元件——模型、資料集、優化器、損失函數，以及蒸餾損失本身——皆以宣告式 YAML 檔案描述。該庫讀取此檔案，建構物件並執行實驗。

主要功能（如 README 所述）

功能	重要性
模組化蒸餾方法 – 實作多種先進的 KD 技術（如 FitNets、Attention Transfer、Relational KD、Variational Information Distillation 等），您只需幾行設定即可嘗試。
前向鈎子管理器 – 無需修改模型的 `forward` 方法即可擷取任意層的中間激活。在蒸餾（教師-學生特徵匹配）與模型分析中極為實用。
「零程式碼」實驗 – 僅透過編輯單一 YAML 檔案即可定義整個流程（資料集、模型、訓練超參數）。README 中甚至展示僅用 YAML 建立 `torchvision.datasets.CIFAR10` 物件的 CIFAR-10 設定範例。
廣泛的任務覆蓋 – 提供影像分類、物件檢測、語意分割，以及 NLP（透過 Hugging-Face Transformers 實作 GLUE 任務）的範例程式碼。
預訓練模型 – 包含部分重新實作的 CIFAR-10/100 模型，並連結至 Hugging-Face Model Hub 上托管的 Transformer 檢查點。
PyTorch 生態系統成員 – 官方列於 PyTorch 生態系統，表示遵循相同的套件與文件規範。
安裝簡單 – `pip install torchdistill`（或透過 `pipenv`）。

如何使用

撰寫一個 YAML 檔案，宣告教師與/或學生模型、資料集、優化器，以及要套用的蒸餾損失。您也可指定要鈎取以進行特徵提取的層。
執行提供的 CLI（或匯入套件）——框架將建構物件、註冊前向鈎子並開始訓練。
可選地，透過 ForwardHookManager.pop_io_dict() 檢視保存的中間張量，用於除錯或研究分析。

典型工作流程範例（CIFAR-10）

models:
  teacher_model:
    key: 'resnet34'
    kwargs:
      pretrained: true
  student_model:
    key: 'resnet18'
    kwargs:
      pretrained: false

datasets:
  cifar10/train: !import_call
    key: 'torchvision.datasets.CIFAR10'
    init:
      kwargs:
        root: '~/datasets/cifar10'
        train: true
        download: true
        transform: !import_call
          key: 'torchvision.transforms.Compose'
          init:
            kwargs:
              transforms:
                - !import_call {key: 'torchvision.transforms.RandomCrop', init: {kwargs: {size: 32, padding: 4}}}
                - !import_call {key: 'torchvision.transforms.RandomHorizontalFlip', init: {kwargs: {p: 0.5}}}
                - !import_call {key: 'torchvision.transforms.ToTensor'}
                - !import_call {key: 'torchvision.transforms.Normalize', init: {kwargs: {mean: [0.49,0.48,0.44], std: [0.24,0.24,0.26]}}}

training:
  epochs: 200
  optimizer: !import_call {key: 'torch.optim.SGD', init: {kwargs: {lr: 0.1, momentum: 0.9, weight_decay: 5e-4}}}
  distillation:
    loss: !import_call {key: 'torchdistill.loss.kd.KDLoss', init: {kwargs: {temperature: 4.0, alpha: 0.7}}}

執行此實驗將使用教師模型的軟化 logits（KD 損失）與標準交叉熵損失來訓練學生的 ResNet-18 模型。

進一步了解

完整 API 文件：https://yoshitomo-matsubara.net/torchdistill/
演示筆記本（例如提取中間表示）位於 demo/ 資料夾中，可直接在 Google Colab 中開啟。
基準測試與範例結果列於專案網站與 examples/ 目錄中。

引用若在論文中使用 torchdistill，請引用兩篇列出的會議論文（2021 年 torchdistill 工作坊論文與 2023 年 torchdistill meets Hugging Face 論文）。README 提供 BibTeX 條目。

總結 – torchdistill 是一個真實、持續維護的知識蒸餾研究庫。它抽象掉了重複的訓練程式碼，支援廣泛的 KD 方法，並與視覺與 NLP 模型皆相容，使任何希望在不深入底層 PyTorch 迴圈的情況下進行教師-學生訓練實驗的人都能受益。

yoshitomo-matsubara/torchdistill

torchdistill – 以設定驅動的知識蒸餾框架

相關