Training with Trainer

Trainer runs the loop for you

Trainer wraps the standard PyTorch training loop: batching, moving tensors to the device, forward pass, loss, backward pass, optimiser and scheduler steps, gradient accumulation, mixed precision, logging, evaluation and checkpointing. You supply a model, datasets, a data collator and a TrainingArguments object. The model computes the loss itself when the batch contains a labels column, so you rarely write a loss function. Key arguments: learning_rate (around 2e-5 to 5e-5 is a common range for fine-tuning BERT-sized encoders), per_device_train_batch_size, num_train_epochs, weight_decay, eval_strategy and save_strategy (for example "epoch"), and bf16=True or fp16=True for mixed precision on supported GPUs. DataCollatorWithPadding pads each batch only to its own longest example. Pass the tokenizer as processing_class so it is saved with checkpoints; older code used the tokenizer= argument, and older releases called eval_strategy evaluation_strategy.

Fine-tuning a sentiment classifier

train_ds and val_ds are tokenized datasets with a labels column, as built in the datasets section.

from transformers import (AutoTokenizer, AutoModelForSequenceClassification,
                          DataCollatorWithPadding, Trainer, TrainingArguments)

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

args = TrainingArguments(
    output_dir="./runs/sentiment",
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=32,
    num_train_epochs=3,
    weight_decay=0.01,
    eval_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
    metric_for_best_model="f1",
    logging_steps=50,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_ds,
    eval_dataset=val_ds,
    data_collator=DataCollatorWithPadding(tok),
    processing_class=tok,
    compute_metrics=compute_metrics,  # defined in the next topic
)
trainer.train()
trainer.save_model("./runs/sentiment/best")

Out of GPU memory? Accumulate gradients

If a batch of 16 does not fit, use per_device_train_batch_size=4 with gradient_accumulation_steps=4. The optimiser still steps once per 16 examples, so training behaves much like the larger batch while using a quarter of the activation memory.

त्वरित जाँच: Your batches contain a `labels` column. Where does the loss come from when you use Trainer with a model from AutoModelForSequenceClassification?

• You must always write a custom loss function
• The data collator computes it
• TrainingArguments computes it from metric_for_best_model
• The model's forward pass computes it and returns it in the outputs