torchdistill.core

torchdistill.core.forward_hook

torchdistill.core.forward_hook.get_device_index(data)[source]

Gets device index of tensor in given data.

Parameters:: data (torch.Tensor or abc.Mapping or tuple or list) – tensor or data structure containing tensor.
Returns:: device index.
Return type:: int or str or None

torchdistill.core.forward_hook.clear_io_dict_values(io_dict)[source]

Clears the values stored in an I/O dict, leaving an empty dict for each module path.

Forward hooks repopulate the I/O type entries at the next forward pass, so this does not affect the registered forward hooks. This is the shared implementation behind ForwardHookManager.clear_io_dict() and torchdistill.core.util.clear_io_dict().

Parameters:: io_dict (dict) – I/O dict whose stored values should be cleared.

torchdistill.core.forward_hook.register_forward_hook_with_dict(root_module, module_path, requires_input, requires_output, io_dict, accumulates=False)[source]

Registers a forward hook for a child module to store its input and/or output in io_dict.

Parameters:

root_module (nn.Module) – root module (e.g., model).
module_path (str) – path to target child module.
requires_input (bool) – if True, stores input to the target child module.
requires_output (bool) – if True, stores output from the target child module.
io_dict (dict) – dict to store the target child module’s input and/or output.
accumulates (bool) – if True, appends input/output across forward passes instead of overwriting. Useful for autoregressive generation where the same module is called multiple times.

Returns:

removable forward hook handle.

Return type:

torch.utils.hook.RemovableHandle

class torchdistill.core.forward_hook.ForwardHookManager(target_device)[source]

A forward hook manager for PyTorch modules.

Parameters:: target_device (torch.device or str) – target device.

Example:

>>> import torch
>>> from torchvision import models
>>> from torchdistill.core.forward_hook import ForwardHookManager
>>> device = torch.device('cpu')
>>> forward_hook_manager = ForwardHookManager(device)
>>> model = models.resnet18()
>>> forward_hook_manager.add_hook(model, 'layer2')
>>> x = torch.rand(16, 3, 224, 224)
>>> y = model(x)
>>> io_dict = forward_hook_manager.pop_io_dict()
>>> layer2_input_tensor = io_dict['layer2']['input']
>>> layer2_output_tensor = io_dict['layer2']['output']

add_hook(root_module, module_path, requires_input=True, requires_output=True, accumulates=False, stacks_accumulated=False)[source]

Registers a forward hook for a child module to store its input and/or output.

Parameters:

root_module (nn.Module) – root module (e.g., model).
module_path (str) – path to target child module.
requires_input (bool) – if True, stores input to the target child module.
requires_output (bool) – if True, stores output from the target child module.
accumulates (bool) – if True, appends input/output across forward passes instead of overwriting. Useful for autoregressive generation where the same module is called multiple times.
stacks_accumulated (bool) – if True, stacks the accumulated per-step tensors into a single torch.Tensor via torch.stack when pop_io_dict is called. Requires accumulates=True and that all per-step tensors have the same shape.

Raises:

ValueError – if stacks_accumulated=True but accumulates=False.

add_hook_to_module(target_module, module_path, requires_input=True, requires_output=True, accumulates=False, stacks_accumulated=False)[source]

Registers a forward hook for an already resolved target module.

Use this instead of add_hook() when the target module cannot be reached by a module path from a single root module e.g., when the module was added to a redesigned model and thus should be looked up in a different module tree than the other target modules.

Parameters:

target_module (nn.Module) – target module to be hooked.
module_path (str) – module path used as a key of the I/O dict.
requires_input (bool) – if True, stores input to the target module.
requires_output (bool) – if True, stores output from the target module.
accumulates (bool) – if True, appends input/output across forward passes instead of overwriting. Useful for autoregressive generation where the same module is called multiple times.
stacks_accumulated (bool) – if True, stacks the accumulated per-step tensors into a single torch.Tensor via torch.stack when pop_io_dict() is called. Requires accumulates=True and that all per-step tensors have the same shape.

Raises:

ValueError – if stacks_accumulated=True but accumulates=False.

Returns:

pair of module path and removable forward hook handle.

Return type:

(str, torch.utils.hook.RemovableHandle)

pop_io_dict()[source]

Pops I/O dict after gathering tensors on self.target_device.

For module paths registered with accumulates=True, the returned value per I/O type is a list of tensors/outputs (one per forward pass) instead of a single tensor. The list is not stacked into a torch.Tensor because tensor shapes may vary across steps (e.g., the sequence dimension grows in models without KV-cache). If the shapes are uniform, callers can stack manually: torch.stack(io_dict[module_path]['output']), or register the hook with stacks_accumulated=True to have this done automatically.

Returns:: I/O dict that contains input and/or output tensors with a module path as a key.
Return type:: dict

pop_io_dict_from_device(device)[source]

Pops I/O dict for a specified device.

Parameters:: device (torch.device) – device to pop I/O dict.
Returns:: I/O dict that contains input and/or output tensors with a module path as a key.
Return type:: dict

change_target_device(target_device)[source]

Updates the target device with a new target_device.

Parameters:: target_device (torch.device or str) – new target device.

clear_io_dict()[source]

Clears the values stored in the I/O dict, keeping the registered forward hooks.

Each module path is left with an empty dict, and the forward hooks repopulate the I/O type entries at the next forward pass. Use clear() instead if the registered forward hooks should be unregistered as well.

clear()[source]: Clears I/O dict and forward hooks registered in the instance.

torchdistill.core.interfaces

torchdistill.core.interfaces.forward_proc

torchdistill.core.interfaces.forward_proc.forward_all(model, *args, **kwargs)[source]

Performs forward computation using *args and **kwargs.

Parameters:

model (nn.Module) – model.
args (tuple) – variable-length arguments for forward.
kwargs (dict) – kwargs for forward.

Returns:

model’s forward output.

Return type:

Any

torchdistill.core.interfaces.forward_proc.forward_batch_only(model, sample_batch, targets=None, supp_dict=None, **kwargs)[source]

Performs forward computation using sample_batch only.

Parameters:

model (nn.Module) – model.
sample_batch (Any) – sample batch.
targets (Any) – training targets (won’t be passed to forward).
supp_dict (dict) – supplementary dict (won’t be passed to forward).

Returns:

model’s forward output.

Return type:

Any

torchdistill.core.interfaces.forward_proc.forward_batch_only_as_kwargs(model, sample_batch, targets=None, supp_dict=None)[source]

Performs forward computation using sample_batch only.

Parameters:

model (nn.Module) – model.
sample_batch (dict) – sample batch.
targets (Any) – training targets (won’t be passed to forward).
supp_dict (dict) – supplementary dict (won’t be passed to forward).

Returns:

model’s forward output.

Return type:

Any

torchdistill.core.interfaces.forward_proc.forward_batch_target(model, sample_batch, targets, supp_dict=None, **kwargs)[source]

Performs forward computation using sample_batch and targets only.

Parameters:

model (nn.Module) – model.
sample_batch (Any) – sample batch.
targets (Any) – training targets.
supp_dict (dict) – supplementary dict (won’t be passed to forward).

Returns:

model’s forward output.

Return type:

Any

torchdistill.core.interfaces.forward_proc.forward_batch_supp_dict(model, sample_batch, targets, supp_dict=None, **kwargs)[source]

Performs forward computation using sample_batch and supp_dict only.

Parameters:

model (nn.Module) – model.
sample_batch (Any) – sample batch.
targets (Any) – training targets (won’t be passed to forward).
supp_dict (dict) – supplementary dict.

Returns:

model’s forward output.

Return type:

Any

torchdistill.core.interfaces.forward_proc.forward_batch4sskd(model, sample_batch, targets=None, supp_dict=None, **kwargs)[source]

Performs forward computation using sample_batch only for the SSKD method.

Guodong Xu, Ziwei Liu, Xiaoxiao Li, Chen Change Loy: “Knowledge Distillation Meets Self-Supervision” @ ECCV 2020 (2020)

Parameters:

model (nn.Module) – model.
sample_batch (Any) – sample batch.
targets (Any) – training targets (won’t be passed to forward).
supp_dict (dict) – supplementary dict (won’t be passed to forward).

Returns:

model’s forward output.

Return type:

Any

torchdistill.core.interfaces.pre_epoch_proc

torchdistill.core.interfaces.pre_epoch_proc.default_pre_epoch_process_with_teacher(self, epoch=None, **kwargs)[source]

Performs pre-epoch process for distillation box.

Parameters:

self (torchdistill.core.distillation.DistillationBox or torchdistill.core.training.TrainingBox) – distillation box.
epoch (int) – epoch for DistributedSampler.

torchdistill.core.interfaces.pre_epoch_proc.default_pre_epoch_process_without_teacher(self, epoch=None, **kwargs)[source]

Performs pre-epoch process for training box.

Parameters:

self (torchdistill.core.distillation.DistillationBox or torchdistill.core.training.TrainingBox) – distillation box.
epoch (int) – epoch for DistributedSampler.

torchdistill.core.interfaces.pre_forward_proc

torchdistill.core.interfaces.post_forward_proc

torchdistill.core.interfaces.post_forward_proc.default_post_forward_process(self, loss, metrics=None, **kwargs)[source]

Performs post-forward process for distillation box.

Parameters:

self (torchdistill.core.distillation.DistillationBox or torchdistill.core.training.TrainingBox) – distillation box.
loss (torch.Tensor) – loss tensor.
metrics (Any) – metric for ReduceLROnPlateau.step.

torchdistill.core.interfaces.post_epoch_proc

torchdistill.core.interfaces.post_epoch_proc.default_post_epoch_process_with_teacher(self, metrics=None, **kwargs)[source]

Performs post-epoch process for distillation box.

Parameters:

self (torchdistill.core.distillation.DistillationBox) – distillation box.
metrics (Any) – metric for ReduceLROnPlateau.step.

torchdistill.core.interfaces.post_epoch_proc.default_post_epoch_process_without_teacher(self, metrics=None, **kwargs)[source]

Performs post-epoch process for training box.

Parameters:

self (torchdistill.core.training.TrainingBox) – training box.
metrics (Any) – metric for ReduceLROnPlateau.step.

torchdistill.core.interfaces.registry

torchdistill.core.interfaces.registry.register_pre_epoch_proc_func(arg=None, **kwargs)[source]

Registers a pre-epoch process function for torchdistill.core.distillation.DistillationBox and torchdistill.core.training.TrainingBox.

Parameters:: arg (Callable or None) – function to be registered as a pre-epoch process function.
Returns:: registered pre-epoch process function.
Return type:: Callable

Note

The function will be registered as an option of the pre-epoch process function. You can choose the registered function by specifying the name of the function or key you used for the registration, in a training configuration used for torchdistill.core.distillation.DistillationBox or torchdistill.core.training.TrainingBox.

If you want to register the function with a key of your choice, add key to the decorator as below:

>>> from torchdistill.core.interfaces.registry import register_pre_epoch_proc_func
>>> @register_pre_epoch_proc_func(key='my_custom_pre_epoch_proc_func')
>>> def new_pre_epoch_proc(self, epoch=None, **kwargs):
>>>     print('This is my custom pre-epoch process function')

In the example, new_pre_epoch_proc function is registered with a key “my_custom_pre_epoch_proc_func”. When you configure torchdistill.core.distillation.DistillationBox or torchdistill.core.training.TrainingBox, you can choose the new_pre_epoch_proc function by “my_custom_pre_epoch_proc_func”.

torchdistill.core.interfaces.registry.register_pre_forward_proc_func(arg=None, **kwargs)[source]

Registers a pre-forward process function for torchdistill.core.distillation.DistillationBox and torchdistill.core.training.TrainingBox.

Parameters:: arg (Callable or None) – function to be registered as a pre-forward process function.
Returns:: registered pre-forward process function.
Return type:: Callable

Note

The function will be registered as an option of the pre-forward process function. You can choose the registered function by specifying the name of the function or key you used for the registration, in a training configuration used for torchdistill.core.distillation.DistillationBox or torchdistill.core.training.TrainingBox.

If you want to register the function with a key of your choice, add key to the decorator as below:

>>> from torchdistill.core.interfaces.registry import register_pre_forward_proc_func
>>> @register_pre_forward_proc_func(key='my_custom_pre_forward_proc_func')
>>> def new_pre_forward_proc(self, *args, **kwargs):
>>>     print('This is my custom pre-forward process function')

In the example, new_pre_forward_proc function is registered with a key “my_custom_pre_forward_proc_func”. When you configure torchdistill.core.distillation.DistillationBox or torchdistill.core.training.TrainingBox, you can choose the new_pre_forward_proc function by “my_custom_pre_forward_proc_func”.

torchdistill.core.interfaces.registry.register_forward_proc_func(arg=None, **kwargs)[source]

Registers a forward process function for torchdistill.core.distillation.DistillationBox and torchdistill.core.training.TrainingBox.

Parameters:: arg (Callable or None) – function to be registered as a forward process function.
Returns:: registered forward process function.
Return type:: Callable

Note

The function will be registered as an option of the forward process function. You can choose the registered function by specifying the name of the function or key you used for the registration, in a training configuration used for torchdistill.core.distillation.DistillationBox or torchdistill.core.training.TrainingBox.

If you want to register the function with a key of your choice, add key to the decorator as below:

>>> from torchdistill.core.interfaces.registry import register_forward_proc_func
>>> @register_forward_proc_func(key='my_custom_forward_proc_func')
>>> def new_forward_proc(model, sample_batch, targets=None, supp_dict=None, **kwargs):
>>>     print('This is my custom forward process function')

In the example, new_forward_proc function is registered with a key “my_custom_forward_proc_func”. When you configure torchdistill.core.distillation.DistillationBox or torchdistill.core.training.TrainingBox, you can choose the new_forward_proc function by “my_custom_forward_proc_func”.

torchdistill.core.interfaces.registry.register_post_forward_proc_func(arg=None, **kwargs)[source]

Registers a post-forward process function for torchdistill.core.distillation.DistillationBox and torchdistill.core.training.TrainingBox.

Parameters:: arg (Callable or None) – function to be registered as a post-forward process function.
Returns:: registered post-forward process function.
Return type:: Callable

Note

The function will be registered as an option of the post-forward process function. You can choose the registered function by specifying the name of the function or key you used for the registration, in a training configuration used for torchdistill.core.distillation.DistillationBox or torchdistill.core.training.TrainingBox.

If you want to register the function with a key of your choice, add key to the decorator as below:

>>> from torchdistill.core.interfaces.registry import register_post_forward_proc_func
>>> @register_post_forward_proc_func(key='my_custom_post_forward_proc_func')
>>> def new_post_forward_proc(self, loss, metrics=None, **kwargs):
>>>     print('This is my custom post-forward process function')

In the example, new_post_forward_proc function is registered with a key “my_custom_post_forward_proc_func”. When you configure torchdistill.core.distillation.DistillationBox or torchdistill.core.training.TrainingBox, you can choose the new_post_forward_proc function by “my_custom_post_forward_proc_func”.

torchdistill.core.interfaces.registry.register_post_epoch_proc_func(arg=None, **kwargs)[source]

Registers a post-epoch process function for torchdistill.core.distillation.DistillationBox and torchdistill.core.training.TrainingBox.

Parameters:: arg (Callable or None) – function to be registered as a post-epoch process function.
Returns:: registered post-epoch process function.
Return type:: Callable

Note

The function will be registered as an option of the post-epoch process function. You can choose the registered function by specifying the name of the function or key you used for the registration, in a training configuration used for torchdistill.core.distillation.DistillationBox or torchdistill.core.training.TrainingBox.

If you want to register the function with a key of your choice, add key to the decorator as below:

>>> from torchdistill.core.interfaces.registry import register_post_epoch_proc_func
>>> @register_post_epoch_proc_func(key='my_custom_post_epoch_proc_func')
>>> def new_post_epoch_proc(self, metrics=None, **kwargs):
>>>     print('This is my custom post-epoch process function')

In the example, new_post_epoch_proc function is registered with a key “my_custom_post_epoch_proc_func”. When you configure torchdistill.core.distillation.DistillationBox or torchdistill.core.training.TrainingBox, you can choose the new_post_epoch_proc function by “my_custom_post_epoch_proc_func”.

torchdistill.core.interfaces.registry.get_pre_epoch_proc_func(key)[source]

Gets a registered pre-epoch process function.

Parameters:: key (str) – unique key to identify the registered pre-epoch process function.
Returns:: registered pre-epoch process function.
Return type:: Callable

torchdistill.core.interfaces.registry.get_pre_forward_proc_func(key)[source]

Gets a registered pre-forward process function.

Parameters:: key (str) – unique key to identify the registered pre-forward process function.
Returns:: registered pre-forward process function.
Return type:: Callable

torchdistill.core.interfaces.registry.get_forward_proc_func(key)[source]

Gets a registered forward process function.

Parameters:: key (str) – unique key to identify the registered forward process function.
Returns:: registered forward process function.
Return type:: Callable

torchdistill.core.interfaces.registry.get_post_forward_proc_func(key)[source]

Gets a registered post-forward process function.

Parameters:: key (str) – unique key to identify the registered post-forward process function.
Returns:: registered post-forward process function.
Return type:: Callable

torchdistill.core.interfaces.registry.get_post_epoch_proc_func(key)[source]

Gets a registered post-epoch process function.

Parameters:: key (str) – unique key to identify the registered post-epoch process function.
Returns:: registered post-epoch process function.
Return type:: Callable

torchdistill.core.interfaces.registry.build_proc_func(proc_config, get_registered_proc_func)[source]

Builds a process function from a process configuration.

The configuration is either a unique key (str) of a registered process function, or a dict with key and optional kwargs. The kwargs are used as default keyword arguments of the process function, and keyword arguments given at call time take precedence over them.

Parameters:

proc_config (str or dict) – process configuration.
get_registered_proc_func (Callable) – getter of a registered process function e.g., get_pre_forward_proc_func().

Returns:

process function.

Return type:

Callable

torchdistill.core.training

class torchdistill.core.training.TrainingBox(model, dataset_dict, train_config, device, device_ids, distributed, lr_factor, accelerator=None)[source]

A single-stage training framework.

Parameters:

model (nn.Module) – model.
dataset_dict (dict) – dict that contains datasets with IDs of your choice.
train_config (dict) – training configuration.
device (torch.device) – target device.
device_ids (list[int]) – target device IDs.
distributed (bool) – whether to be in distributed training mode.
lr_factor (float or int) – multiplier for learning rate.
accelerator (accelerate.Accelerator or None) – Hugging Face accelerator.

setup_data_loaders(train_config)[source]

Sets up training and validation data loaders for the current training stage. This method will be internally called when instantiating this class and when calling MultiStagesTrainingBox.advance_to_next_stage().

Parameters:: train_config (dict) – training configuration.

setup_model(model_config)[source]

Sets up a model for the current training stage. This method will be internally called when instantiating this class and when calling MultiStagesTrainingBox.advance_to_next_stage().

Parameters:: model_config (dict) – model configuration.

setup_loss(train_config)[source]

Sets up a training loss module for the current training stage. This method will be internally called when instantiating this class and when calling MultiStagesTrainingBox.advance_to_next_stage().

Parameters:: train_config (dict) – training configuration.

setup_pre_post_processes(train_config)[source]

Sets up pre/post-epoch/forward processes for the current training stage. This method will be internally called when instantiating this class and when calling MultiStagesTrainingBox.advance_to_next_stage().

Parameters:: train_config (dict) – training configuration.

setup(train_config)[source]

Configures a TrainingBox/MultiStagesTrainingBox for the current training stage. This method will be internally called when instantiating this class and when calling MultiStagesTrainingBox.advance_to_next_stage().

Parameters:: train_config (dict) – training configuration.

property model_io_dict

I/O dict of the model, populated by the forward hooks registered with forward_hook configuration.

Returns:: model I/O dict.
Return type:: dict

property target_model_pairs

Pairs of module path and removable forward hook handle registered for the model.

Returns:: list of pairs of module path and removable forward hook handle.
Return type:: list[(str, torch.utils.hook.RemovableHandle)]

pre_epoch_process(*args, **kwargs)[source]

Performs a pre-epoch process Shows the summary of results.

This should be overridden by all subclasses or defined through setup_pre_post_processes().

pre_forward_process(*args, **kwargs)[source]

Performs a pre-forward process Shows the summary of results.

This should be overridden by all subclasses or defined through setup_pre_post_processes().

forward_process(sample_batch, targets=None, supp_dict=None, **kwargs)[source]

Performs forward computations for a model.

Parameters:

sample_batch (Any) – sample batch.
targets (Any) – training targets.
supp_dict (dict) – supplementary dict.

Returns:

loss tensor.

Return type:

torch.Tensor

post_forward_process(*args, **kwargs)[source]

Performs a post-forward process.

This should be overridden by all subclasses or defined through setup_pre_post_processes().

post_epoch_process(*args, **kwargs)[source]

Performs a post-epoch process.

This should be overridden by all subclasses or defined through setup_pre_post_processes().

clean_modules()[source]: Unfreezes all the modules, clears an I/O dict, unregisters forward hook handles, and clears the handle lists.

class torchdistill.core.training.MultiStagesTrainingBox(model, dataset_dict, train_config, device, device_ids, distributed, lr_factor, accelerator=None)[source]

A multi-stage training framework. This is a subclass of TrainingBox.

Parameters:

model (nn.Module) – model.
dataset_dict (dict) – dict that contains datasets with IDs of your choice.
train_config (dict) – training configuration.
device (torch.device) – target device.
device_ids (list[int]) – target device IDs.
distributed (bool) – whether to be in distributed training mode.
lr_factor (float or int) – multiplier for learning rate.
accelerator (accelerate.Accelerator or None) – Hugging Face accelerator.

save_stage_ckpt(model, local_model_config)[source]

Saves the checkpoint of model for the current training stage.

Parameters:

model (nn.Module) – model to be saved.
local_model_config (dict) – model configuration at the current training stage.

advance_to_next_stage()[source]: Reads the next training stage’s configuration in train_config and advances to the next training stage.

post_epoch_process(*args, **kwargs)[source]

Performs a post-epoch process.

The superclass’s post_epoch_process should be overridden by all subclasses or defined through TrainingBox.setup_pre_post_processes().

torchdistill.core.training.get_training_box(model, dataset_dict, train_config, device, device_ids, distributed, lr_factor, accelerator=None)[source]

Gets a training box.

Parameters:

model (nn.Module) – model.
dataset_dict (dict) – dict that contains datasets with IDs of your choice.
train_config (dict) – training configuration.
device (torch.device) – target device.
device_ids (list[int]) – target device IDs.
distributed (bool) – whether to be in distributed training mode.
lr_factor (float or int) – multiplier for learning rate.
accelerator (accelerate.Accelerator or None) – Hugging Face accelerator.

Returns:

training box.

Return type:

TrainingBox or MultiStagesTrainingBox

torchdistill.core.distillation

class torchdistill.core.distillation.DistillationBox(teacher_model, student_model, dataset_dict, train_config, device, device_ids, distributed, lr_factor, accelerator=None)[source]

A single-stage knowledge distillation framework.

Parameters:

teacher_model (nn.Module) – teacher model.
student_model (nn.Module) – student model.
dataset_dict (dict) – dict that contains datasets with IDs of your choice.
train_config (dict) – training configuration.
device (torch.device) – target device.
device_ids (list[int]) – target device IDs.
distributed (bool) – whether to be in distributed training mode.
lr_factor (float or int) – multiplier for learning rate.
accelerator (accelerate.Accelerator or None) – Hugging Face accelerator.

setup_data_loaders(train_config)[source]

Sets up training and validation data loaders for the current training stage. This method will be internally called when instantiating this class and when calling MultiStagesDistillationBox.advance_to_next_stage().

Parameters:: train_config (dict) – training configuration.

setup_teacher_student_models(teacher_config, student_config)[source]

Sets up teacher and student models for the current training stage. This method will be internally called when instantiating this class and when calling MultiStagesDistillationBox.advance_to_next_stage().

Parameters:

teacher_config (dict) – teacher configuration.
student_config (dict) – student configuration.

setup_loss(train_config)[source]

Sets up a training loss module for the current training stage. This method will be internally called when instantiating this class and when calling MultiStagesDistillationBox.advance_to_next_stage().

Parameters:: train_config (dict) – training configuration.

setup_pre_post_processes(train_config)[source]

Sets up pre/post-epoch/forward processes for the current training stage. This method will be internally called when instantiating this class and when calling MultiStagesDistillationBox.advance_to_next_stage().

Parameters:: train_config (dict) – training configuration.

setup(train_config)[source]

Configures a DistillationBox/MultiStagesDistillationBox for the current training stage. This method will be internally called when instantiating this class and when calling MultiStagesDistillationBox.advance_to_next_stage().

Parameters:: train_config (dict) – training configuration.

property teacher_io_dict

I/O dict of the teacher model, populated by the forward hooks registered with forward_hook configuration.

Returns:: teacher model I/O dict.
Return type:: dict

property student_io_dict

I/O dict of the student model, populated by the forward hooks registered with forward_hook configuration.

Returns:: student model I/O dict.
Return type:: dict

property target_teacher_pairs

Pairs of module path and removable forward hook handle registered for the teacher model.

Returns:: list of pairs of module path and removable forward hook handle.
Return type:: list[(str, torch.utils.hook.RemovableHandle)]

property target_student_pairs

Pairs of module path and removable forward hook handle registered for the student model.

Returns:: list of pairs of module path and removable forward hook handle.
Return type:: list[(str, torch.utils.hook.RemovableHandle)]

pre_epoch_process(*args, **kwargs)[source]

Performs a pre-epoch process Shows the summary of results.

This should be overridden by all subclasses or defined through setup_pre_post_processes().

pre_forward_process(*args, **kwargs)[source]

Performs a pre-forward process Shows the summary of results.

This should be overridden by all subclasses or defined through setup_pre_post_processes().

get_teacher_output(sample_batch, targets, supp_dict, **kwargs)[source]

Gets teacher model’s output.

Parameters:

sample_batch (Any) – sample batch.
targets (Any) – training targets.
supp_dict (dict) – supplementary dict.

Returns:

teacher’s outputs and teacher’s I/O dict.

Return type:

(Any, dict)

forward_process(sample_batch, targets=None, supp_dict=None, **kwargs)[source]

Performs forward computations for teacher and student models.

Parameters:

sample_batch (Any) – sample batch.
targets (Any) – training targets.
supp_dict (dict) – supplementary dict.

Returns:

loss tensor.

Return type:

torch.Tensor

post_forward_process(*args, **kwargs)[source]

Performs a post-forward process.

This should be overridden by all subclasses or defined through setup_pre_post_processes().

post_epoch_process(*args, **kwargs)[source]

Performs a post-epoch process.

This should be overridden by all subclasses or defined through setup_pre_post_processes().

clean_modules()[source]: Unfreezes all the teacher and student modules, clears I/O dicts, unregisters forward hook handles, and clears the handle lists.

class torchdistill.core.distillation.MultiStagesDistillationBox(teacher_model, student_model, dataset_dict, train_config, device, device_ids, distributed, lr_factor, accelerator=None)[source]

A multi-stage knowledge distillation framework. This is a subclass of DistillationBox.

Parameters:

teacher_model (nn.Module) – teacher model.
student_model (nn.Module) – student model.
dataset_dict (dict) – dict that contains datasets with IDs of your choice.
train_config (dict) – training configuration.
device (torch.device) – target device.
device_ids (list[int]) – target device IDs.
distributed (bool) – whether to be in distributed training mode.
lr_factor (float or int) – multiplier for learning rate.
accelerator (accelerate.Accelerator or None) – Hugging Face accelerator.

save_stage_ckpt(model, local_model_config)[source]

Saves the checkpoint of model for the current training stage.

Parameters:

model (nn.Module) – model to be saved.
local_model_config (dict) – model configuration at the current training stage.

advance_to_next_stage()[source]: Reads the next training stage’s configuration in train_config and advances to the next training stage.

post_epoch_process(*args, **kwargs)[source]

Performs a post-epoch process.

The superclass’s post_epoch_process should be overridden by all subclasses or defined through DistillationBox.setup_pre_post_processes().

torchdistill.core.distillation.get_distillation_box(teacher_model, student_model, dataset_dict, train_config, device, device_ids, distributed, lr_factor, accelerator=None)[source]

Gets a distillation box.

Parameters:

teacher_model (nn.Module) – teacher model.
student_model (nn.Module) – student model.
dataset_dict (dict) – dict that contains datasets with IDs of your choice.
train_config (dict) – training configuration.
device (torch.device) – target device.
device_ids (list[int]) – target device IDs.
distributed (bool) – whether to be in distributed training mode.
lr_factor (float or int) – multiplier for learning rate.
accelerator (accelerate.Accelerator or None) – Hugging Face accelerator.

Returns:

distillation box.

Return type:

DistillationBox or MultiStagesDistillationBox

torchdistill.core.util

torchdistill.core.util.get_proc_config(train_config, key)[source]

Extracts a pre/post-epoch/forward process configuration, supporting the deprecated key (e.g., pre_forward_process for pre_forward_proc) for backward compatibility.

Parameters:

train_config (dict) – training configuration.
key (str) – process configuration key e.g., pre_forward_proc.

Returns:

process configuration if either the key or its deprecated version is available, None otherwise.

Return type:

dict or str or None

torchdistill.core.util.add_kwargs_to_io_dict(io_dict, module_path, **kwargs)[source]

Adds kwargs to an I/O dict.

Deprecated since version 1.2.0: Forward hooks initialize their own entries in an I/O dict, and this function is no longer used internally. It will be removed in a future release.

Parameters:

io_dict (dict) – I/O dict.
module_path (str) – module path.
kwargs (dict) – kwargs to be stored in io_dict.

torchdistill.core.util.set_hooks(model, unwrapped_org_model, model_config, forward_hook_manager)[source]

Sets forward hooks for target modules in model, using forward_hook_manager.

model_config['forward_hook'] accepts the following keys:

input: list of module paths whose input should be stored.
output: list of module paths whose output should be stored.
accumulates: bool (applied to all the target modules) or list of module paths whose input/output should be accumulated across forward passes instead of being overwritten. Useful for autoregressive use cases such as on-policy distillation, where the target modules are called once per generated token.
stacks_accumulated: bool or list of module paths whose accumulated per-step tensors should be stacked into a single tensor by torchdistill.core.forward_hook.ForwardHookManager.pop_io_dict(). Requires the same module paths to be accumulated, and that the per-step tensors share the same shape.

Parameters:

model (nn.Module) – model.
unwrapped_org_model (nn.Module) – unwrapped original model.
model_config (dict) – model configuration.
forward_hook_manager (torchdistill.core.forward_hook.ForwardHookManager) – forward hook manager to register the forward hooks with.

Returns:

list of pairs of module path and removable forward hook handle.

Return type:

list[(str, torch.utils.hook.RemovableHandle)]

torchdistill.core.util.wrap_model(model, model_config, device, device_ids=None, distributed=False, find_unused_parameters=False, any_updatable=True)[source]

Wraps model with DataParallel, DistributedDataParallel, FullyShardedDataParallel (FSDP), or FSDP2 (fully_shard) if specified.

model_config['wrapper']['key'] selects the wrapper and accepts one of ‘DataParallel’, ‘DistributedDataParallel’, ‘FullyShardedDataParallel’ (FSDP), or ‘FullyShardedDataParallel2’ (FSDP2). model_config['wrapper']['kwargs'] are forwarded to the wrapper’s constructor (e.g., auto_wrap_policy, sharding_strategy, mixed_precision for FSDP).

Note

FSDP/FSDP2 shard parameters across the process group, so model’s state_dict() will only reflect the local shard. Use torchdistill.common.module_util.get_full_state_dict() / torchdistill.common.module_util.load_full_state_dict() for checkpointing instead of calling state_dict() / load_state_dict() directly.

Parameters:

model (nn.Module) – model.
model_config (dict) – model configuration.
device (torch.device) – target device.
device_ids (list[int]) – target device IDs.
distributed (bool) – whether to be in distributed training mode.
find_unused_parameters (bool) – find_unused_parameters for DistributedDataParallel.
any_updatable (bool) – True if model contains any updatable parameters.

Returns:

wrapped model (or model if wrapper is not specified).

Return type:

nn.Module

torchdistill.core.util.clear_io_dict(model_io_dict)[source]

Clears a model I/O dict’s sub dict(s).

Each module path is left with an empty dict, and the forward hooks repopulate the I/O type entries at the next forward pass.

Note

If you hold a torchdistill.core.forward_hook.ForwardHookManager, prefer its clear_io_dict() method. Both share torchdistill.core.forward_hook.clear_io_dict_values() as their implementation.

Parameters:: model_io_dict (dict) – model I/O dict.

torchdistill.core.util.extract_io_dict(model_io_dict, target_device)[source]

Extracts I/O dict, gathering tensors on target_device.

Deprecated since version 1.2.0: Use torchdistill.core.forward_hook.ForwardHookManager.pop_io_dict() instead, which additionally supports accumulated I/O. Unlike pop_io_dict(), this function always adds a torchdistill.common.constant.SELF_MODULE_PATH entry, so replace io_dict[SELF_MODULE_PATH]['output'] = outputs with io_dict[SELF_MODULE_PATH] = {'output': outputs} when migrating. This function will be removed in a future release.

Parameters:

model_io_dict (dict) – model I/O dict.
target_device (torch.device or str) – target device.

Returns:

extracted I/O dict.

Return type:

dict

torchdistill.core.util.update_io_dict(main_io_dict, sub_io_dict)[source]

Updates an I/O dict with a sub I/O dict.

Parameters:

main_io_dict (dict) – main I/O dict to be updated.
sub_io_dict (dict) – sub I/O dict.