11.4 Using Transformer Models in PyTorch

Two Practical Routes: Torch Modules vs Hugging Face

In PyTorch you will commonly use Transformer models in two ways. The first is to use the Transformer building blocks that ship with PyTorch, such as nn.TransformerEncoderLayer and nn.TransformerEncoder, when you want a lightweight model and full control over the architecture. The second is to use a pretrained Transformer from the Hugging Face Transformers library when you want strong out of the box performance and a standard training and inference interface. The underlying ideas are the same, you pass token representations through stacked self attention blocks, but the surrounding details, especially tokenization and pretrained weights, differ.

This chapter focuses on the practical mechanics of using Transformer models in code, including the key input shapes, masks, and the minimal forward pass patterns you will rely on later for fine tuning and batching.

Using PyTorch’s `nn.TransformerEncoder` for Custom Models

PyTorch provides general purpose Transformer components. For most beginner friendly NLP setups, the encoder only part is enough to start, for example for text classification.

A typical custom encoder pipeline has three steps. First, map token ids to vectors using an embedding layer. Second, add positional information. Third, run the sequence through a Transformer encoder stack and pool the result into a single vector for classification or regression.

In modern PyTorch, many sequence modules support batch_first=True, which is usually the simplest choice. When batch_first=True, the main sequence tensor shape is (batch, seq_len, d_model).

A minimal model skeleton looks like this:

import math
import torch
import torch.nn as nn
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=512, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float32).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2, dtype=torch.float32) * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer("pe", pe.unsqueeze(0))  # (1, max_len, d_model)
    def forward(self, x):
        # x: (batch, seq_len, d_model)
        x = x + self.pe[:, :x.size(1)]
        return self.dropout(x)
class TransformerClassifier(nn.Module):
    def __init__(self, vocab_size, d_model=128, nhead=4, num_layers=2, num_classes=2, max_len=512):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)
        self.pos = PositionalEncoding(d_model, max_len=max_len)
        layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=4*d_model,
            dropout=0.1,
            batch_first=True
        )
        self.encoder = nn.TransformerEncoder(layer, num_layers=num_layers)
        self.classifier = nn.Linear(d_model, num_classes)
    def forward(self, input_ids, padding_mask=None):
        # input_ids: (batch, seq_len)
        x = self.embed(input_ids)              # (batch, seq_len, d_model)
        x = self.pos(x)                        # (batch, seq_len, d_model)
        x = self.encoder(x, src_key_padding_mask=padding_mask)  # same shape
        pooled = x[:, 0]                       # simple choice if you reserve token 0 as a [CLS]-like token
        logits = self.classifier(pooled)       # (batch, num_classes)
        return logits

The main thing you must supply correctly is the padding mask.

Padding Masks: The Most Common Source of Bugs

When you batch sequences of different lengths, you pad them to the same seq_len. Self attention must not treat padding as real tokens. In PyTorch encoder modules, src_key_padding_mask is typically a boolean tensor of shape (batch, seq_len) where True means “this position is padding, ignore it.”

If you store padded token id 0 as your padding id, then the mask is often:

padding_mask = (input_ids == 0)  # True at pads

Be careful because some libraries and some functions use the opposite convention.

If you get shape errors, check whether batch_first=True is set. If it is not, PyTorch expects (seq_len, batch, d_model).

Causal Masks: Preventing Attention to the Future

If you are doing autoregressive modeling, you must prevent each position from attending to later positions. That is done with a causal attention mask, often an upper triangular matrix. In PyTorch you pass this as mask or src_mask depending on the module.

A common causal mask shape is (seq_len, seq_len). A typical construction is:

seq_len = input_ids.size(1)
causal_mask = torch.triu(torch.ones(seq_len, seq_len, device=input_ids.device), diagonal=1).bool()

You generally use either a padding mask, a causal mask, or both. For encoder only classification tasks, you usually only need the padding mask.

Pooling the Encoder Output for Classification

A Transformer encoder outputs a vector per token. For classification you need one vector for the whole sequence. Common pooling choices are to take the first token, to take the mean over non padding tokens, or to use attention pooling. The simplest pooling that works well in practice is mean pooling with a padding mask.

Mean pooling with masking looks like this:

# x: (batch, seq_len, d_model)
# padding_mask: (batch, seq_len) True at pads
valid = (~padding_mask).unsqueeze(-1)         # (batch, seq_len, 1)
x_masked = x * valid                          # pads become zero vectors
lengths = valid.sum(dim=1).clamp(min=1)       # (batch, 1)
pooled = x_masked.sum(dim=1) / lengths        # (batch, d_model)

If you instead rely on a dedicated first token, you must ensure your data pipeline always places a special token there and that the model can learn to use it.

Using Pretrained Transformers with Hugging Face in PyTorch

For most real world beginner projects, you will use pretrained models. Hugging Face Transformers provides PyTorch modules that already include embeddings, positional encodings, and Transformer blocks, plus a tokenizer that converts text into input_ids and attention_mask.

A minimal inference example looks like this:

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
name = "distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(name)
model = AutoModelForSequenceClassification.from_pretrained(name, num_labels=2)
texts = ["this movie was great", "this was boring"]
batch = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
with torch.no_grad():
    out = model(**batch)
logits = out.logits
pred = logits.argmax(dim=-1)

In this setup, batch is a dictionary. For many models it contains input_ids and attention_mask. The attention mask is usually 1 for real tokens and 0 for padding.

You do not manually create positional encodings or attention masks beyond providing attention_mask and possibly token_type_ids for some architectures.

Moving Transformer Batches to Devices Correctly

When using Hugging Face, remember that the tokenizer output is a dictionary of tensors. Moving only input_ids to the GPU is a common mistake. Move the whole batch:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
batch = {k: v.to(device) for k, v in batch.items()}

This pattern also works for custom models when you keep inputs in a dict.

What the Forward Pass Returns

PyTorch custom modules return whatever you define, usually logits or predictions. Hugging Face models return a structured output object with attributes like logits, and sometimes loss if you pass labels.

For sequence classification, if you pass labels shaped (batch,) for class indices, you can get loss directly:

labels = torch.tensor([1, 0], device=device)
out = model(**batch, labels=labels)
loss = out.loss
logits = out.logits

This is convenient because the model knows which loss to apply for the head you selected.

Common Shape Expectations You Should Memorize

Token ids are integer tensors shaped (batch, seq_len). Attention masks in Hugging Face are shaped (batch, seq_len) with 1s and 0s. Encoder outputs are typically (batch, seq_len, hidden_size) and classification logits are (batch, num_labels).

Minimal Checklist Before You Train

Before running a full training loop, verify that a single forward pass works and that the output shapes match your task. For custom PyTorch Transformer encoders, confirm batch_first usage, confirm src_key_padding_mask boolean polarity, and confirm your pooling method ignores padding. For Hugging Face models, confirm you move every tensor in the tokenizer batch to the same device as the model and that num_labels matches your dataset.

These mechanics are the core of “using Transformer models in PyTorch,” and they set you up for the next chapters where you fine tune models, manage tokenizers, and build correct batches.