torchdistill

The torchdistill repository is associated with several research projects and papers focusing on knowledge distillation, split computing, and neural feature compression. Key areas include:

• Knowledge Distillation: Research into the role of projectors, normalization, and soft maximum functions in distillation efficiency.
• Split Computing (SC2): Benchmarking and developing supervised compression methods for resource-constrained edge computing.
• Neural Feature Compression: Methods like FOOL and FrankenSplit for addressing downlink bottlenecks in satellite and mobile edge computing.
• NLP Integration: Using torchdistill to harmonize with Hugging Face libraries for reproducible NLP studies (e.g., GLUE benchmark).

The torchdistill.losses package provides a hierarchical structure for implementing and managing loss functions used in knowledge distillation. The package is organized into the following functional layers:

• torchdistill.losses.registry: Manages the registration and retrieval of loss functions.
• torchdistill.losses.high_level: Contains high-level loss abstractions, typically used directly in distillation training loops.
• torchdistill.losses.mid_level: Contains mid-level loss components that may serve as building blocks for high-level losses.
• torchdistill.losses.util: Provides utility functions and helpers for loss computation.

The torchdistill.core package provides the fundamental building blocks for knowledge distillation experiments. It is organized into several functional sub-modules:

• forward_hook: Tools for managing and registering forward hooks to extract intermediate layer activations.
• interfaces: Defines the standard interfaces (abstract base classes or protocols) for various lifecycle hooks used during training and distillation.
• training: Core logic for managing the training loop.
• distillation: Logic specifically related to the distillation process (e.g., loss functions, teacher-student interactions).
• util: General utility functions used across the core package.

Beyond standard classification models, torchdistill.models provides several specialized sub-packages:

• torchdistill.models.official: Likely contains official implementations of specific architectures.
• torchdistill.models.adaptation: Models designed for adaptation tasks (e.g., domain adaptation or task-specific fine-tuning).
• torchdistill.models.wrapper: Utilities for wrapping existing models to modify their behavior or interface for distillation purposes.

PAD-L2 is a two-stage training method. To use it:

1. Stage 1: Train a model using the L2 (CE + L2) method.
2. Stage 2: Rename the cse_l2/ directory in your unzipped checkpoint directory to ce_l2/.
3. Stage 2 Execution: Load the checkpoint file designated in the pad_l2 YAML file and run the training command.

When running distributed training for PAD-L2, ensure the batch_size in train_data_loader is adjusted so that batch_size * NUM_GPUS = 512.

In torchdistill, experiments are designed by editing declarative PyYAML configuration files rather than writing Python code. This approach allows you to:

• Perform Knowledge Distillation: Implement state-of-the-art distillation methods by configuring teacher and student models.
• Extract Intermediate Representations: Instead of reimplementing models to access internal layers, you can specify the module path(s) directly in the YAML file to extract features.
• Conduct General Deep Learning Studies: You can perform standard training (without teachers) by simply excluding teacher entries from your configuration file.

Sample configurations can be found in the configs/sample/ directory of the official repository.

To extend the training or distillation process, you can implement custom logic by adhering to the interfaces defined in torchdistill.core.interfaces. These interfaces allow you to inject code at specific stages of the execution lifecycle:

• forward_proc: Logic applied during the forward pass.
• pre_epoch_proc: Logic executed before each epoch starts.
• pre_forward_proc: Logic executed immediately before the forward pass.
• post_forward_proc: Logic executed immediately after the forward pass.
• post_epoch_proc: Logic executed after each epoch completes.

The repository provides a wide variety of declarative PyYAML configuration files that implement different knowledge distillation (KD) research papers. These samples can be used as templates for your own experiments.

Key categories of available configurations include:

• Classic KD Methods: e.g., kd (Hinton et al.), fitnet (Romero et al.).
• Advanced Representation/Relational KD: e.g., rkd (Relational KD), crd (Contrastive Representation Distillation), at (Attention Transfer).
• Task-Specific KD: Configurations for ilsvrc2012 (ImageNet) and coco2017 (Object Detection/Segmentation).
• Recent Research: Includes implementations for papers from CVPR 2024, 2026, and TPAMI 2025.

To train DeepLabv3 models without using distributed processes, ensure that no checkpoint files exist at the dst_ckpt path specified in the student_model section of your YAML configuration. This prevents the script from loading an existing checkpoint and forces training from scratch.

Use the --run_log flag to specify a log file for the training process.

# For DeepLabv3 with ResNet-50
python3 examples/torchvision/semantic_segmentation.py \
    --config configs/official/pascal_voc2012/yoshitomo-matsubara/nlp-oss2023/deeplabv3_resnet50.yaml \
    --run_log log/deeplabv3_resnet50.log 

# For DeepLabv3 with ResNet-101
python3 examples/torchvision/semantic_segmentation.py \
    --config configs/official/pascal_voc2012/yoshitomo-matsubara/nlp-oss2023/deeplabv3_resnet101.yaml \
    --run_log log/deeplabv3_resnet101.log