Trainer overview

Trainer is the batteries-included training loop for PyTorch Transformers models.

Adapted for this playbook from the 🤗 Transformers documentation by Hugging Face. Official page: https://huggingface.co/docs/transformers/trainer. Images © Hugging Face (documentation-images) unless noted. This is not a substitute for the upstream docs — verify against the current version.

Covers Hugging Face pages: trainer, trainer_customize, trainer_callbacks, data_collators, optimizers, hpo_train, trainer_recipes

Executive summary

Use Trainer when you want a standard loop with logging, eval, checkpointing, mixed precision and distributed hooks. Subclass or use callbacks when you need custom behaviour.

Minimal shape

from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments
from datasets import load_dataset

model_id = "bert-base-uncased"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSequenceClassification.from_pretrained(model_id, num_labels=2)

# prepare tokenized datasets ...
args = TrainingArguments(
    output_dir="./out",
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    num_train_epochs=3,
    eval_strategy="epoch",
    fp16=True,
)
trainer = Trainer(model=model, args=args, train_dataset=train_ds, eval_dataset=eval_ds, processing_class=tok)
trainer.train()

Customisation surfaces

Hook	Doc
Subclass methods	trainer_customize
Callbacks	trainer_callbacks
Data collators	data_collators
Optimisers / schedulers	optimizers
Hyperparameter search	hpo_train

Source: Hugging Face Transformers documentation.

Key takeaways

• TrainingArguments is the control plane for Trainer.
• Prefer callbacks over forking the whole loop.

What to do next

Fine-tuning.