peal.training.trainers.distill_predictor

peal.training.trainers.distill_predictor(predictor_distillation, base_path, predictor, predictor_datasource, replace_with_activation=None, tracking_level=4, predictor_distilled=None, only_last_layer=False, continue_training=False, task_config=None)[source]

Train a student predictor on data relabelled by predictor.

The relabelling strategy is chosen from the datasource: with distill_from == "dataset" the original labels are kept; a (DataloaderMixer, WeightedDataloaderList) pair is relabelled with distill_dataloader_mixer() / distill_binary_dataset(); Image2MixedDataset data with distill_binary_dataset(); and Image2ClassDataset data with distill_1ofn_dataset(). The student is built from predictor_distillation.architecture, or is a deep copy of the teacher, and is trained by a ModelTrainer into <base_path>/distilled_predictor.

Parameters:

  • predictor_distillation (str, dict or PredictorConfig) – Student config; distill_from and architecture are read here.

  • base_path (str) – Directory for relabelled data and the trained student.

  • predictor (torch.nn.Module) – Teacher whose predictions become the labels.

  • predictor_datasource (sequence) – (train, validation) datasets or loaders of the teacher.

  • replace_with_activation ({"leakysoftplus", "leakyrelu", None}, optional) – Swap the student’s ReLUs (leakysoftplus only for a freshly built student, leakyrelu only for a given predictor_distilled).

  • tracking_level (int, optional) – Verbosity written into the student config. Default 4.

  • predictor_distilled (torch.nn.Module, optional) – Student to train instead of building one.

  • only_last_layer (bool, optional) – Forwarded to ModelTrainer.

  • continue_training (bool, optional) – Forwarded to ModelTrainer.fit().

  • task_config (TaskConfig, optional) – Overrides the task config of the relabelled datasets.

Returns:

The trained student.

Return type:

torch.nn.Module

Raises:

Exception – When the datasource type matches no relabelling strategy.