yoshitomo-matsubara/torchdistill

torchdistill – A configuration‑driven framework for knowledge distillation

What it is

• An open‑source Python library built on top of PyTorch that lets you run knowledge‑distillation experiments (teacher‑student model training) without writing custom training loops.
• All components – models, datasets, optimizers, loss functions, and the distillation losses themselves – are described in a declarative YAML file. The library reads the file, builds the objects, and runs the experiment.

Key features (as described in the README)

Feature	Why it matters
Modular distillation methods – implements many state‑of‑the‑art KD techniques (e.g., FitNets, Attention Transfer, Relational KD, Variational Information Distillation, etc.) so you can try them out with a few lines of config.
Forward‑hook manager – lets you capture intermediate activations from any layer without changing the model’s forward method. Useful for both distillation (teacher‑student feature matching) and model analysis.
“No‑code” experiments – by editing a single YAML file you can define the whole pipeline (datasets, models, training hyper‑parameters). The README even shows a CIFAR‑10 config that creates torchvision.datasets.CIFAR10 objects entirely from YAML.
Broad task coverage – example scripts are provided for image classification, object detection, semantic segmentation, and NLP (GLUE tasks via Hugging‑Face Transformers).
Pre‑trained models – includes some re‑implemented CIFAR‑10/100 models and links to transformer checkpoints hosted on the Hugging‑Face Model Hub.
PyTorch Ecosystem member – officially listed in the PyTorch ecosystem, meaning it follows the same packaging and documentation conventions.
Easy installation – pip install torchdistill (or via pipenv).

How you would use it

1. Write a YAML file that declares the teacher and/or student models, the dataset, optimizer, and which distillation loss to apply. You can also specify which layers to hook for feature extraction.
2. Run the provided CLI (or import the library) – the framework builds the objects, registers the forward hooks, and starts training.
3. Optionally, inspect the saved intermediate tensors via ForwardHookManager.pop_io_dict() for debugging or research analysis.

Typical workflow example (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}}}

Running the experiment will train the student ResNet‑18 using the teacher’s softened logits (KD loss) plus the standard cross‑entropy loss.

Where to learn more

• Full API docs: https://yoshitomo-matsubara.net/torchdistill/
• Demo notebooks (e.g., extracting intermediate representations) are in the demo/ folder and can be opened directly in Google Colab.
• Benchmarks and sample results are listed on the project site and in the examples/ directory.

Citation If you use torchdistill in a paper, cite the two listed conference papers (2021 torchdistill workshop paper and 2023 torchdistill meets Hugging Face paper). The README provides the BibTeX entries.

---

Bottom line – torchdistill is a genuine, actively maintained library for reproducible knowledge‑distillation research. It abstracts away boilerplate training code, supports a wide range of KD methods, and works with both vision and NLP models, making it useful for anyone who wants to experiment with teacher‑student training without digging into low‑level PyTorch loops.