torchdistill: knowledge distillation experiments defined in YAML, not Python

The problem torchdistill targets: distillation code that nobody can rerun

Knowledge distillation papers are usually accompanied by training scripts that hard-code one teacher, one student, one loss and one dataset split. Reproducing a result means reading someone else's training loop closely enough to find where the temperature, the alpha weighting and the intermediate feature alignment actually happen. torchdistill's answer is to make the experiment itself the artifact. The README states that you design experiments by editing a declarative YAML config file instead of Python code, and that the same framework covers general deep learning experiments without a teacher, simply by excluding the teacher entries from the config. The audience is therefore narrow and specific: researchers and engineers who need to run, compare or extend published distillation methods, and who value a config file they can diff against a paper's appendix over a training script they have to reverse engineer. The repository describes 26 knowledge distillation methods presented at venues including TPAMI, CVPR, ICLR, ECCV, NeurIPS, ICCV and AAAI, with trained models, training logs and configurations published alongside them.

How the YAML config actually instantiates a PyTorch model

The mechanism is a custom YAML tag. Config values that look like ordinary mappings are wrapped in !import_call, which carries a key naming a fully qualified Python object and an init block describing how to construct it. The framework resolves the key, calls the object with the given arguments, and nests the result. The README's example builds a CIFAR-10 training dataset by importing torchvision.datasets.CIFAR10 and passing it a transform that is itself an !import_call to torchvision.transforms.Compose, whose transforms list contains further !import_call entries for RandomCrop, RandomHorizontalFlip, ToTensor and Normalize. YAML anchors are used to share values: the root directory is defined once and referenced with *root_dir, and the normalization mean and std are captured under &normalize_kwargs so the test transform can reuse them.

The second mechanism is the ForwardHookManager, which is what makes distillation between differently shaped models tractable. Instead of rewriting a teacher's forward method to return intermediate tensors, you register hooks on named submodules and read the captured inputs and outputs afterwards. The README's example registers a hook on conv1 requesting only its input, on layer1.0.bn2 requesting both input and output, and on fc requesting only its output, then calls pop_io_dict() to retrieve a dictionary keyed by module name. That dictionary is what a distillation loss consumes. The design keeps model code untouched, which matters because a modified forward signature is exactly the kind of change that makes a published checkpoint unusable.

Installing torchdistill and running a first experiment

The package is published on PyPI as torchdistill, and pyproject.toml declares that it requires Python 3.10 or newer. Install it into a virtual environment with pip:

pip install torchdistill

That pulls the runtime dependencies listed in pyproject.toml, which include torch>=2.12.0, torchvision>=0.27.0, numpy, pyyaml>=6.0, scipy and cython. Check that the resolved PyTorch satisfies the floor before you download any dataset, because a mismatch here surfaces much later as a confusing import error.

The fastest way to see the config mechanism work is the README's own snippet, which loads a YAML file and pulls instantiated datasets out of it by name. The convention is that the top-level datasets mapping is keyed by strings like cifar10/train and cifar10/test, so the config file is both the recipe and the lookup table:

from torchdistill.common import yaml_util
config = yaml_util.load_yaml_file('./test.yaml')
train_dataset = config['datasets']['cifar10/train']
test_dataset = config['datasets']['cifar10/test']

To run a complete experiment rather than a fragment, the repository ships configs under configs/ (with sample configs in configs/sample/) and notebooks under demo/. The demos cover CIFAR training and distillation, intermediate representation extraction, and GLUE fine-tuning with and without distillation. If you want to inspect intermediate tensors before committing to a distillation run, the ForwardHookManager example in the README is the shortest path:

import torch
from torchvision import models
from torchdistill.core.forward_hook import ForwardHookManager

model = models.resnet18(pretrained=False)
device = torch.device('cpu')
forward_hook_manager = ForwardHookManager(device)
forward_hook_manager.add_hook(model, 'layer1.0.bn2', requires_input=True, requires_output=True)
io_dict = forward_hook_manager.pop_io_dict()

What you should see is a dictionary keyed by the module path you registered, with input and output entries present only for the flags you set to True.

Where the config-driven approach costs you

The declarative layer is a real abstraction, and abstractions leak. Anything the framework has not abstracted has to be expressed as an !import_call to your own code, at which point you are maintaining both a Python module and a YAML schema for it. The README does not document a rollback or migration path between config formats, and there is no statement about backward compatibility of configs across releases, so a config written against an older version is not guaranteed to load unchanged. The release notes for v1.1.4 record the end of Python 3.9 support, which is a concrete example of the framework's own surface moving.

The dependency floor is the second constraint. torch>=2.12.0 and torchvision>=0.27.0 are not soft preferences; they are the versions the project declares. If your cluster image, your vendor's CUDA build or your other research code pins an older PyTorch, torchdistill is the wrong tool until that pin moves. The third case is methodological: the repository implements a fixed catalogue of 26 methods. If your idea is a distillation loss that is not in that catalogue, you will be adding it, and the value of the config layer drops sharply once most of your experiment lives in a custom module. Finally, the README is explicit that citations should point to the associated papers rather than the GitHub repository, which tells you the project positions itself as an implementation of published work, not as a general-purpose training library.

torchdistill, RepDistiller and DistillKit: three different bets

The searches that lead people here are for RepDistiller and DistillKit, which are the natural comparison points, and the difference is where each puts the experiment definition. RepDistiller is organized around a common training script with a method selector, so adding a distillation method means editing shared training code. torchdistill pushes the definition out into a per-experiment YAML file and keeps the training code generic, which is why it can also run non-distillation experiments by dropping the teacher entries. DistillKit is a collection of distillation implementations, closer in spirit to a library you call than to a framework that runs your experiment for you. The consequence for you: if you want to call a distillation loss from inside your own training loop, a library-shaped project fits better; if you want the experiment itself to be a reviewable file, torchdistill's model is the one that matches. The README's framing of a config as a summary of your experiment is the clearest statement of that bet.

Maintenance, licensing and what a version bump costs

The repository is not archived, and the last push was on 2026-08-28. The most recent release, v1.1.5 on 2026-08-05, added a knowledge distillation method, FSDP/FSDP2 support and experiment tracking. Before that, v1.1.4 on 2025-12-24 added a method and bug fixes and ended Python 3.9 support, and v1.1.3 on 2025-05-11 added a text classification example plus interface and YAML utility updates. The pattern is a few releases a year, each adding methods or infrastructure rather than reshaping the config format, which is roughly what you want if you are pinning a version for a paper.

The licence is MIT, declared in LICENSE and referenced from pyproject.toml via license = { file = "LICENSE" }. MIT is permissive, so incorporating the framework into internal or commercial work is generally straightforward, but the repository does not state a policy on the licences of the datasets and pretrained checkpoints the configs reference, and those are separate works with their own terms. Check each dataset and checkpoint you pull in rather than assuming the framework's licence covers them. This is a description of what the repository declares, not legal advice.

Editorial conclusion

Adopt torchdistill if you are reproducing a published distillation method or comparing several of them and want the experiment recorded as a config file rather than scattered training scripts. Do not adopt it if you need a distillation method the repository does not implement, or if you cannot move to Python 3.10 and a recent PyTorch. Before committing, install it in a clean environment and confirm the dependency floor resolves, then check that the method you intend to use has a config under configs/ and a corresponding entry in the repository's citation list, because those two artifacts are the evidence that the implementation is maintained rather than merely present.

Frequently asked questions

Official sources

1. License: MIT
2. Project website
3. README
4. Releases
5. yoshitomo-matsubara/torchdistill on GitHub