yoshitomo-matsubara/torchdistill

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）。

如何使用

1. 编写一个 YAML 文件，声明教师和/或学生模型、数据集、优化器以及要应用的蒸馏损失。您还可以指定要钩取以进行特征提取的层。
2. 运行提供的 CLI（或导入库）——框架将构建对象、注册前向钩子并开始训练。
3. 可选地，通过 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 循环的情况下进行教师-学生训练实验的人都能受益。