🤗 Transformers provides a Trainer class optimized for training 🤗 Transformers models, making it easier to start training without manually writing your own training loop. The Trainer API supports a wide range of training options and features such as logging, gradient accumulation, and mixed precision.

Start by loading your model and specify the number of expected labels. From the Yelp Review dataset card, you know there are five labels:

>>> from transformers import AutoModelForSequenceClassification

>>> model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=5)

Training hyperparameters

Next, create a TrainingArguments class which contains all the hyperparameters you can tune as well as flags for activating different training options. For this tutorial you can start with the default training hyperparameters, but feel free to experiment with these to find your optimal settings.

Specify where to save the checkpoints from your training:

>>> from transformers import TrainingArguments

>>> training_args = TrainingArguments(output_dir="test_trainer")

Metrics

Trainer does not automatically evaluate model performance during training. You will need to pass Trainer a function to compute and report metrics. The 🤗 Datasets library provides a simple accuracy function you can load with the load_metric (see this tutorial for more information) function:

>>> import numpy as np
>>> from datasets import load_metric

>>> metric = load_metric("accuracy")

Call compute on metric to calculate the accuracy of your predictions. Before passing your predictions to compute, you need to convert the predictions to logits (remember all 🤗 Transformers models return logits):

>>> def compute_metrics(eval_pred):
...     logits, labels = eval_pred
...     predictions = np.argmax(logits, axis=-1)
...     return metric.compute(predictions=predictions, references=labels)

If you’d like to monitor your evaluation metrics during fine-tuning, specify the evaluation_strategy parameter in your training arguments to report the evaluation metric at the end of each epoch:

>>> from transformers import TrainingArguments, Trainer

>>> training_args = TrainingArguments(output_dir="test_trainer", evaluation_strategy="epoch")

Trainer

Create a Trainer object with your model, training arguments, training and test datasets, and evaluation function:

>>> trainer = Trainer(
...     model=model,
...     args=training_args,
...     train_dataset=small_train_dataset,
...     eval_dataset=small_eval_dataset,
...     compute_metrics=compute_metrics,
... )

Then fine-tune your model by calling train():

>>> trainer.train()

🤗 Transformers models also supports training in TensorFlow with the Keras API.

Convert dataset to TensorFlow format

The DefaultDataCollator assembles tensors into a batch for the model to train on. Make sure you specify return_tensors to return TensorFlow tensors:

>>> from transformers import DefaultDataCollator

>>> data_collator = DefaultDataCollator(return_tensors="tf")

Next, convert the tokenized datasets to TensorFlow datasets with the to_tf_dataset method. Specify your inputs in columns, and your label in label_cols:

>>> tf_train_dataset = small_train_dataset.to_tf_dataset(
...     columns=["attention_mask", "input_ids", "token_type_ids"],
...     label_cols=["labels"],
...     shuffle=True,
...     collate_fn=data_collator,
...     batch_size=8,
... )

>>> tf_validation_dataset = small_eval_dataset.to_tf_dataset(
...     columns=["attention_mask", "input_ids", "token_type_ids"],
...     label_cols=["labels"],
...     shuffle=False,
...     collate_fn=data_collator,
...     batch_size=8,
... )

Compile and fit

Load a TensorFlow model with the expected number of labels:

>>> import tensorflow as tf
>>> from transformers import TFAutoModelForSequenceClassification

>>> model = TFAutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=5)

Then compile and fine-tune your model with fit as you would with any other Keras model:

>>> model.compile(
...     optimizer=tf.keras.optimizers.Adam(learning_rate=5e-5),
...     loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
...     metrics=tf.metrics.SparseCategoricalAccuracy(),
... )

>>> model.fit(tf_train_dataset, validation_data=tf_validation_dataset, epochs=3)

Trainer takes care of the training loop and allows you to fine-tune a model in a single line of code. For users who prefer to write their own training loop, you can also fine-tune a 🤗 Transformers model in native PyTorch.

At this point, you may need to restart your notebook or execute the following code to free some memory:

del model
del pytorch_model
del trainer
torch.cuda.empty_cache()

Next, manually postprocess tokenized_dataset to prepare it for training.

1. Remove the text column because the model does not accept raw text as an input:

   Copied

   >>> tokenized_datasets = tokenized_datasets.remove_columns(["text"])

2. Rename the label column to labels because the model expects the argument to be named labels:

   Copied

   >>> tokenized_datasets = tokenized_datasets.rename_column("label", "labels")

3. Set the format of the dataset to return PyTorch tensors instead of lists:

   Copied

   >>> tokenized_datasets.set_format("torch")

Then create a smaller subset of the dataset as previously shown to speed up the fine-tuning:

>>> small_train_dataset = tokenized_datasets["train"].shuffle(seed=42).select(range(1000))
>>> small_eval_dataset = tokenized_datasets["test"].shuffle(seed=42).select(range(1000))

DataLoader

Create a DataLoader for your training and test datasets so you can iterate over batches of data:

>>> from torch.utils.data import DataLoader

>>> train_dataloader = DataLoader(small_train_dataset, shuffle=True, batch_size=8)
>>> eval_dataloader = DataLoader(small_eval_dataset, batch_size=8)

Load your model with the number of expected labels:

>>> from transformers import AutoModelForSequenceClassification

>>> model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=5)

Optimizer and learning rate scheduler

Create an optimizer and learning rate scheduler to fine-tune the model. Let’s use the AdamW optimizer from PyTorch:

>>> from torch.optim import AdamW

>>> optimizer = AdamW(model.parameters(), lr=5e-5)

Create the default learning rate scheduler from Trainer:

>>> from transformers import get_scheduler

>>> num_epochs = 3
>>> num_training_steps = num_epochs * len(train_dataloader)
>>> lr_scheduler = get_scheduler(
...     name="linear", optimizer=optimizer, num_warmup_steps=0, num_training_steps=num_training_steps
... )

Lastly, specify device to use a GPU if you have access to one. Otherwise, training on a CPU may take several hours instead of a couple of minutes.

>>> import torch

>>> device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
>>> model.to(device)

Great, now you are ready to train! 🥳

Training loop

To keep track of your training progress, use the tqdm library to add a progress bar over the number of training steps:

>>> from tqdm.auto import tqdm

>>> progress_bar = tqdm(range(num_training_steps))

>>> model.train()
>>> for epoch in range(num_epochs):
...     for batch in train_dataloader:
...         batch = {k: v.to(device) for k, v in batch.items()}
...         outputs = model(**batch)
...         loss = outputs.loss
...         loss.backward()

...         optimizer.step()
...         lr_scheduler.step()
...         optimizer.zero_grad()
...         progress_bar.update(1)

Metrics

Just like how you need to add an evaluation function to Trainer, you need to do the same when you write your own training loop. But instead of calculating and reporting the metric at the end of each epoch, this time you will accumulate all the batches with add_batch and calculate the metric at the very end.

>>> metric = load_metric("accuracy")
>>> model.eval()
>>> for batch in eval_dataloader:
...     batch = {k: v.to(device) for k, v in batch.items()}
...     with torch.no_grad():
...         outputs = model(**batch)

...     logits = outputs.logits
...     predictions = torch.argmax(logits, dim=-1)
...     metric.add_batch(predictions=predictions, references=batch["labels"])

>>> metric.compute()