Fine-Tuning DistilBERT for IMDB Sentiment

Fine-tuning continues training a pretrained model on labeled data for a specific task. This example adds a binary classification head to DistilBERT and trains it on IMDB reviews. A validation split is used during development, while the official test set remains untouched until training is complete.

Configuration

DistilBERT[1] is a six-layer encoder. The settings below fix the data-split and initialization seed. A maximum length of 256 reduces memory use by truncating longer reviews.

import time
import torch

from tqdm.auto import tqdm
from torch.optim import AdamW
from torch.utils.data import DataLoader
from datasets import DatasetDict, load_dataset
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    DataCollatorWithPadding,
    set_seed,
)

MODEL_NAME = "distilbert-base-uncased"
BATCH_SIZE = 16
EPOCHS = 3
MAX_LEN = 256
LEARNING_RATE = 2e-5
SEED = 42

set_seed(SEED)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

Data

IMDB contains 25,000 labeled training reviews and 25,000 labeled test reviews[2]. Ten percent of the training split becomes validation data. Padding is deferred until batching so each batch is padded only to its longest sequence[3].

raw = load_dataset("stanfordnlp/imdb")
split = raw["train"].train_test_split(
    test_size=0.1,
    seed=SEED,
    stratify_by_column="label",
)

data = DatasetDict({
    "train": split["train"],
    "validation": split["test"],
    "test": raw["test"],
})

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

def tokenize(batch):
    return tokenizer(
        batch["text"],
        truncation=True,
        max_length=MAX_LEN,
    )

data = data.map(tokenize, batched=True, remove_columns=["text"])
collator = DataCollatorWithPadding(tokenizer=tokenizer)

loaders = {
    name: DataLoader(
        dataset_split,
        batch_size=BATCH_SIZE,
        shuffle=(name == "train"),
        collate_fn=collator,
    )
    for name, dataset_split in data.items()
}

Model and evaluation

The sequence-classification model adds a randomly initialized two-class head to the pretrained encoder. AdamW updates the head and all encoder parameters. Accuracy is easy to interpret here because the dataset is balanced.

model = AutoModelForSequenceClassification.from_pretrained(
    MODEL_NAME,
    num_labels=2,
).to(device)

optimizer = AdamW(model.parameters(), lr=LEARNING_RATE)

@torch.inference_mode()
def accuracy(model, loader):
    model.eval()
    correct = 0
    total = 0

    for batch in loader:
        batch = {key: value.to(device) for key, value in batch.items()}
        labels = batch.pop("labels")
        predictions = model(**batch).logits.argmax(dim=-1)
        correct += (predictions == labels).sum().item()
        total += labels.numel()

    return correct / total

Training

Validation accuracy is reported after each epoch. The test loader is evaluated once after all updates are complete. Runtime and accuracy come from the execution rather than a hard-coded estimate.

started_at = time.perf_counter()

for epoch in range(1, EPOCHS + 1):
    model.train()
    total_loss = 0.0
    examples_seen = 0

    progress = tqdm(loaders["train"], desc=f"Epoch {epoch}/{EPOCHS}")
    for batch in progress:
        batch = {key: value.to(device) for key, value in batch.items()}

        optimizer.zero_grad(set_to_none=True)
        output = model(**batch)
        output.loss.backward()
        optimizer.step()

        batch_size = batch["labels"].size(0)
        total_loss += output.loss.item() * batch_size
        examples_seen += batch_size
        progress.set_postfix(loss=f"{output.loss.item():.4f}")

    val_accuracy = accuracy(model, loaders["validation"])
    print(
        f"epoch={epoch} "
        f"train_loss={total_loss / examples_seen:.4f} "
        f"val_accuracy={val_accuracy:.4f}"
    )

test_accuracy = accuracy(model, loaders["test"])
elapsed_minutes = (time.perf_counter() - started_at) / 60

print(f"test_accuracy={test_accuracy:.4f}")
print(f"elapsed_minutes={elapsed_minutes:.1f}")

Interpreting the output

The validation values describe performance during development. The final test value is the only test measurement. Accuracy and runtime depend on the seed, package versions, hardware, and truncation length, so any published result should include those details alongside the output.