Runners

class compressors.distillation.runners.DistilRunner(output_hidden_states: bool = True, apply_probability_shift: bool = False, *runner_args, **runner_kwargs)

class compressors.distillation.runners.EndToEndDistilRunner(hidden_state_loss: Optional[Union[str, Callable]] = None, logits_diff_loss: Union[str, Callable] = 'kl_loss', loss_weights: Optional[Dict] = None, num_train_teacher_epochs: Optional[int] = None, *runner_args, **runner_kwargs)

End to end runner for distillation. Can be used without additional callbacks.

Parameters

• hidden_state_loss – loss between last hidden states. Defaults to None.

• logits_diff_loss – loss between logits. Defaults to kl_loss.

• loss_weights – weights for losses in final sum.

• num_train_teacher_epochs – If set to None, then just distilling the knowledge. If int then train teacher model for provided number of epochs. Defaults to None.

• *runner_args – runner args

• **runner_kwargs – runner kwargs

get_callbacks(stage: str) → Dict[str, catalyst.core.callback.Callback]

Returns the callbacks for a given stage.

get_stage_len(stage: str) → int

Returns the stage length in epochs for a given stage.

predict_batch(batch: Mapping[str, Any], **kwargs) → Mapping[str, Any]

Run model inference on specified data batch.

Parameters

• batch – dictionary with data batches from DataLoader.

• **kwargs – additional kwargs to pass to the model

Returns

model output dictionary

Return type

Mapping

Raises

NotImplementedError – if not implemented yet

property stages

Experiment’s stage names (array with one value).

class compressors.distillation.runners.HFDistilRunner(*args, **kwargs)

Simple runner for transformer model.