TopolearnReal projects need class weights, special losses or extra logging; Trainer has hooks for each.
Trainer is designed to be extended rather than replaced.
Callbacks observe and steer training. Subclass TrainerCallback and implement events like on_log, on_evaluate, on_epoch_end or on_step_end. They receive args, state (step, epoch, log history, best metric) and control, whose flags such as control.should_stop = True change behaviour. Built-ins include EarlyStoppingCallback(early_stopping_patience=3), which requires load_best_model_at_end=True and a metric_for_best_model. Integration callbacks for TensorBoard, Weights & Biases or MLflow are added via report_to.
Custom loss: subclass Trainer and override compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None), for example to apply class weights. Pop the labels, run the model, compute the loss, and return (loss, outputs) when return_outputs is true. Alternatively pass compute_loss_func= to the constructor.
Custom optimizer and scheduler: pass optimizers=(optimizer, scheduler), or override create_optimizer, e.g. to use different learning rates for the backbone and the head.
Callbacks should not mutate the model or data; keep them for logging and control flow. For anything requiring a fundamentally different loop (reinforcement learning, GANs), look at TRL or write a plain PyTorch loop with Accelerate.
Example
import torch
from transformers import EarlyStoppingCallback, Trainer, TrainerCallback
class WeightedLossTrainer(Trainer):
def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
labels = inputs.pop("labels")
outputs = model(**inputs)
weights = torch.tensor([1.0, 4.0], device=outputs.logits.device) # rare class x4
loss = torch.nn.functional.cross_entropy(outputs.logits, labels, weight=weights)
return (loss, outputs) if return_outputs else loss
class PrintEval(TrainerCallback):
def on_evaluate(self, args, state, control, metrics=None, **kwargs):
print(f"epoch {state.epoch}: {metrics.get('eval_f1')}")
# trainer = WeightedLossTrainer(
# model=model, args=args, # args: load_best_model_at_end=True, metric_for_best_model="f1"
# train_dataset=train, eval_dataset=val, compute_metrics=compute_metrics,
# callbacks=[EarlyStoppingCallback(early_stopping_patience=2), PrintEval()],
# )Remember
- Callbacks observe events and can stop or redirect training via
control - Override
compute_loss(or passcompute_loss_func) for custom losses - EarlyStoppingCallback needs
load_best_model_at_endandmetric_for_best_model
Watch out for
- Forgetting to return (loss, outputs) when return_outputs=True, which breaks evaluation and prediction.