Trainer, and what it saves you writing

The whole thing, for a classifier

import numpy as np, evaluate
from transformers import (AutoModelForSequenceClassification, AutoTokenizer,
                          TrainingArguments, Trainer, DataCollatorWithPadding)

tok = AutoTokenizer.from_pretrained("distilbert-base-uncased")
model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased", num_labels=3,
    id2label={0: "neg", 1: "neu", 2: "pos"},
    label2id={"neg": 0, "neu": 1, "pos": 2},
)

ds = ds.map(lambda b: tok(b["text"], truncation=True, max_length=256), batched=True)

f1 = evaluate.load("f1")
def metrics(p):
    preds = np.argmax(p.predictions, axis=1)
    return f1.compute(predictions=preds, references=p.label_ids, average="macro")

args = TrainingArguments(
    output_dir="out",
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    num_train_epochs=3,
    eval_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
    metric_for_best_model="f1",
    push_to_hub=True,
    hub_model_id="your-name/reviews-classifier",
)

trainer = Trainer(
    model=model, args=args,
    train_dataset=ds["train"], eval_dataset=ds["validation"],
    data_collator=DataCollatorWithPadding(tok),
    compute_metrics=metrics,
)
trainer.train()
trainer.push_to_hub()

That is a complete, publishable fine-tune. Note id2label passed at construction: the thing whose absence produces LABEL_0 outputs forever, and the cheapest fix in this course.

What Trainer is doing for you

It writes, correctly, the loop you would otherwise get subtly wrong: batching and shuffling, moving tensors to the device, the forward and backward passes, gradient accumulation, learning-rate scheduling and warmup, gradient clipping, mixed precision, evaluation at intervals, checkpoint saving and rotation, resumption, logging, and multi-GPU distribution through accelerate.

None of those is hard alone. Getting all of them right together, and keeping them right, is why nobody writes their own loop for ordinary work.

Metrics, and evaluate

pip install evaluate

acc = evaluate.load("accuracy")
f1  = evaluate.load("f1")
wer = evaluate.load("wer")          # speech
rouge = evaluate.load("rouge")      # summarisation

evaluate loads standard metric implementations from the Hub, which means everybody computing F1 is computing the same F1. That matters more than it sounds: average="macro" and average="micro" differ enormously on imbalanced data, and two teams quoting “F1” without saying which have not compared anything.

For a classifier on imbalanced classes, macro-F1 is usually the honest headline, with per-class precision and recall printed beside it.

Resuming, which you will need

trainer.train(resume_from_checkpoint=True)

On free hardware the session will end. With save_strategy="epoch" or save_steps=200 and this one argument, a disconnection costs minutes rather than hours. Combine it with push_to_hub=True and checkpoints leave the machine as well, which is the module-three discipline applied to training.

save_total_limit=2 stops checkpoints filling the disk: each one is a full copy of the model, and three copies of a 7B model is 42 GB.

The evaluation loop, and its cost

eval_strategy="steps" with eval_steps=100 runs the whole validation set every hundred steps. On a large validation set that can take longer than the training between evaluations, and people then conclude that training is slow when most of the wall-clock time is measurement.

Two fixes. Keep the validation set small: a thousand rows is plenty for watching a curve, and the real evaluation happens once at the end on a larger held-out set. And set per_device_eval_batch_size explicitly, because it inherits the training batch size by default and evaluation has different memory behaviour.

eval_dataset=ds["validation"].select(range(1000)),
per_device_eval_batch_size=32,

Seq2seq and generation

from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments

args = Seq2SeqTrainingArguments(..., predict_with_generate=True, generation_max_length=128)

predict_with_generate=True makes evaluation actually generate text rather than scoring forced labels, which is what you need for ROUGE or word error rate. Without it your evaluation metric measures something other than what the model will do at inference, and looks better than reality.

When to write your own loop

When you need a custom loss, an unusual optimisation schedule, or a training procedure with several models in it: reinforcement learning, adversarial setups, some distillation schemes. Subclassing Trainer and overriding compute_loss covers most of these and keeps everything else:

class WeightedTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False, **kw):
        labels = inputs.pop("labels")
        out = model(**inputs)
        loss = torch.nn.functional.cross_entropy(
            out.logits, labels, weight=class_weights.to(out.logits.device))
        return (loss, out) if return_outputs else loss

That is class weighting for imbalanced data, which the data-quality lesson recommended, in eight lines.

Do this now

Train a classifier on one per cent of a dataset for one epoch. It will be bad, and that is not the point: the point is that the whole loop (data, metrics, checkpoints, push) runs end to end in a few minutes. Get the plumbing right on something worthless before you spend hours on something you care about.