distillation — Model Optimizer 0.27.1 (original) (raw)

API for converting a model into a modelopt.torch.distill.DistillationModel to be used directly in training.

Functions

convert(model, mode)

Main conversion function to turn a student model into a distillation-ready model.

Parameters:

• model (Module) – The base model to be used as the student.
• mode (_ModeDescriptor | str | list [ _ModeDescriptor | str ] | list [ tuple [ str , dict [ str , Any ] ] ]) –
  A (list of) string(s) or Mode(s) or a list of tuples containing the mode and its config indicating the desired mode(s) (and configurations) for the convert process. Modes set up the model for different algorithms for model optimization. The following modes are available:
  • "kd_loss": Themodel will be converted into meta-model encapsulating both teacher and student. The mode’s config is described inKDLossConfig.
    If the mode argument is specified as a dictionary, the keys should indicate the mode and the values specify the per-mode configuration.

Returns:

An instance of DistillationModel <modelopt.torch.distill.DistillationModel.

Return type:

Module

export(model)

Export a distillation meta-model to the original student model.

Parameters:

model (Module) – Model to be exported out of a distillation mode into the student only.

Returns:

The inner student model.

Return type:

Module