integrations.kd.trainer
integrations.kd.trainer
KD trainer
Classes
| Name | Description |
|---|---|
| AxolotlKDTrainer | Custom trainer subclass for Knowledge Distillation (KD) |
AxolotlKDTrainer
integrations.kd.trainer.AxolotlKDTrainer(*args, **kwargs)Custom trainer subclass for Knowledge Distillation (KD)
Methods
| Name | Description |
|---|---|
| compute_loss | How the loss is computed by Trainer. By default, all models return the loss in the first element. |
compute_loss
integrations.kd.trainer.AxolotlKDTrainer.compute_loss(
model,
inputs,
return_outputs=False,
num_items_in_batch=None,
)How the loss is computed by Trainer. By default, all models return the loss in the first element.
Subclass and override for custom behavior.