config

Configurations for distillation modes.

Classes

ExportStudentConfig	Configuration for the export_student mode.
KDLossConfig	Configuration for the Knowledge-Distillation mode.
LayerwiseKDConfig	Configuration for the Layerwise Knowledge-Distillation mode.

class ExportStudentConfig

Bases: ModeloptBaseConfig

Configuration for the export_student mode.

This mode is used to export a student model after distillation training/fine-tuning.

model_config = {'extra': 'forbid', 'validate_assignment': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class KDLossConfig

Bases: ModeloptBaseConfig

Configuration for the Knowledge-Distillation mode.

This mode is used to distill knowledge from a teacher model to a student model.

criterion: _Loss | dict[tuple[str, str], _Loss] | None

expose_minimal_state_dict: bool

classmethod format_criterion(criterion)

Ensure criterion is a mapping from layer names to loss (potentially entire module).

Parameters:

criterion (_Loss | dict[tuple[str, str], _Loss] | None)

Return type:

dict[tuple[str, str], _Loss]

loss_balancer: Any | None

model_config = {'arbitrary_types_allowed': True, 'extra': 'forbid', 'validate_assignment': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_dump(*args, **kwargs)

Dump the config to a dictionary but avoid serializing teacher model to dict.

This avoids issues when the teacher is a tuple with callable and args. If any of the args are Dataclasses, they are dumped as a dict and cannot be restored with their class.

Return type:

dict[str, Any]

teacher_model: Module | type[Module] | tuple | Callable | None

class LayerwiseKDConfig

Bases: KDLossConfig

Configuration for the Layerwise Knowledge-Distillation mode.

This mode is used to distill knowledge from a teacher model to a student model using layerwise distillation.

classmethod format_criterion(criterion)

Ensure criterion is a mapping from layer names to loss (potentially entire module).

Parameters:

criterion (_Loss | dict[tuple[str, str], _Loss] | None)

Return type:

dict[tuple[str, str], _Loss]

model_config = {'arbitrary_types_allowed': True, 'extra': 'forbid', 'validate_assignment': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].