Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

RNN Architecture › PyTorch nn.RNN: Code Implementation

Michael BrenndoerferMay 7, 20258 min read

Language AI HandbookAll chapters

Part 11 of 67Recurrent Neural Networks

Chapter 74RNN Architecture › PyTorch nn.RNN: Code Implementation

Reading section 13 of 19 · PyTorch nn.RNN: Code Implementation · 8 minSection 13 of 19 · 8 min

Section 13 of 19 · 8 min read

RNN Architecture › PyTorch nn.RNN: Code ImplementationOne section · Sections in this chapter13 / 19 ↓

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

PyTorch provides nn.RNN as a ready-to-use module. Let us implement a character-level language model from scratch, which will demonstrate all the architectural pieces in a running system.

Setup: Character-Level TokenizationLink Copied

We will work with a small corpus of text and build a vocabulary of individual characters.

In[8]:

Code

# Small training corpus
corpus = (
    "the quick brown fox jumps over the lazy dog "
    "recurrent neural networks process sequences step by step "
    "the hidden state carries context forward through time "
    "language models predict the next word given previous words "
)

# Build character vocabulary
chars = sorted(set(corpus))
vocab_size = len(chars)
char_to_idx = {c: i for i, c in enumerate(chars)}
idx_to_char = {i: c for i, c in enumerate(chars)}

# Encode the corpus as integer indices
encoded = [char_to_idx[c] for c in corpus]

Out[9]:

Console

Corpus length: 214 characters
Vocabulary size: 27 unique characters
Characters: ' abcdefghijklmnopqrstuvwxyz'

A vocabulary of around 30 characters is typical for character-level models: all lowercase letters, spaces, and punctuation. Each character maps to an integer index, and integer indices will be converted to embeddings inside the model.

Character-level modeling is a useful pedagogical setting because the vocabulary is small and tractable, the task (predicting the next character) is well-defined, and you can directly inspect what the model generates. The challenge for the model is that it must learn word boundaries, morphology, and even syntax purely from the character stream, without any explicit tokenization.

Dataset PreparationLink Copied

For language modeling, the input at each step is a token, and the target at each step is the next token. We create overlapping windows of a fixed sequence length.

In[10]:

Code

import torch

seq_len = 30  # Number of characters per training window
batch_size = 16

# Build (input, target) pairs: target is input shifted by one position
sequences = []
targets = []
for i in range(0, len(encoded) - seq_len - 1, 1):
    sequences.append(encoded[i : i + seq_len])
    targets.append(encoded[i + 1 : i + seq_len + 1])

# Convert to tensors
X = torch.tensor(sequences, dtype=torch.long)  # shape: (N, seq_len)
Y = torch.tensor(targets, dtype=torch.long)  # shape: (N, seq_len)

# DataLoader for minibatch training
dataset = torch.utils.data.TensorDataset(X, Y)
loader = torch.utils.data.DataLoader(
    dataset, batch_size=batch_size, shuffle=True
)

Out[11]:

Console

Number of training windows: 183
Input tensor shape: torch.Size([183, 30])
Target tensor shape: torch.Size([183, 30])
Number of batches per epoch: 12

Each training window has seq_len input characters and seq_len target characters, shifted by one position. At every position, the model must predict the next character from the current hidden state. For example, if the input window is "the quick brow", the target window is "he quick brown". The model is simultaneously trained to predict each character given the preceding ones.

The dense supervision from this formulation is important. At seq_len = 30 and batch size 16, each forward pass generates 480 predictions (30 per sequence times 16 sequences), each supervised by a cross-entropy loss. This dense signal is what makes language modeling such an effective training objective.

Defining the RNN Language ModelLink Copied

In[12]:

Code

import torch.nn as nn


class CharRNN(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_size, num_layers=1):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers

        # Embedding layer: integer index -> dense vector
        self.embedding = nn.Embedding(vocab_size, embed_dim)

        # Core RNN: processes the embedded sequence
        self.rnn = nn.RNN(
            input_size=embed_dim,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,  # Input/output shape: (batch, seq, feature)
            nonlinearity="tanh",  # Default; 'relu' is the other option
        )

        # Output projection: hidden state -> vocab logits
        self.fc = nn.Linear(hidden_size, vocab_size)

    def forward(self, x, h=None):
        # x:   (batch, seq_len)   integer token indices
        # h:   (num_layers, batch, hidden_size) optional initial hidden state
        embedded = self.embedding(x)  # -> (batch, seq_len, embed_dim)
        out, h_n = self.rnn(embedded, h)  # -> (batch, seq_len, hidden_size)
        logits = self.fc(out)  # -> (batch, seq_len, vocab_size)
        return logits, h_n

    def init_hidden(self, batch_size, device):
        return torch.zeros(
            self.num_layers, batch_size, self.hidden_size, device=device
        )

The model has three components. The embedding layer converts integer token indices into dense vectors, equivalent to looking up a row in an embedding matrix. The nn.RNN layer processes the embedded sequence and returns both the sequence of hidden states (one per time step) and the final hidden state. The linear layer projects each hidden state to vocabulary logits. With batch_first=True, all tensors have the batch dimension first, which is the conventional layout for training loops.

Notice that the forward method accepts an optional initial hidden state h. During generation, we will pass the hidden state from one call to the next, allowing the model to maintain context across generation steps. During training with independent windows, we pass None and let PyTorch initialize to zeros automatically.

Training Loop with Teacher ForcingLink Copied

In[13]:

Code

import torch
import torch.nn as nn
import torch.optim as optim

# Hyperparameters
embed_dim = 64
hidden_size = 256
num_epochs = 30
lr = 0.003

device = torch.device("cpu")
model = CharRNN(vocab_size, embed_dim, hidden_size).to(device)
optimizer = optim.Adam(model.parameters(), lr=lr)
criterion = nn.CrossEntropyLoss()

train_losses = []

for epoch in range(num_epochs):
    epoch_loss = 0.0
    model.train()

    for x_batch, y_batch in loader:
        x_batch = x_batch.to(device)
        y_batch = y_batch.to(device)

        # Teacher forcing: input = ground-truth tokens at every step
        # The model receives the correct previous token, not its own prediction
        h = model.init_hidden(x_batch.size(0), device)
        logits, _ = model(x_batch, h)  # logits: (batch, seq_len, vocab_size)

        # Reshape for cross-entropy: (batch*seq_len, vocab_size) vs (batch*seq_len,)
        loss = criterion(logits.view(-1, vocab_size), y_batch.view(-1))

        optimizer.zero_grad()
        loss.backward()
        # Gradient clipping to prevent exploding gradients
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0)
        optimizer.step()

        epoch_loss += loss.item()

    avg_loss = epoch_loss / len(loader)
    train_losses.append(avg_loss)

Out[14]:

Console

Initial training loss (epoch 1):  2.1569
Final training loss (epoch 30): 0.0919
Loss reduction: 95.7%

The teacher forcing pattern is evident in the training loop: x_batch contains the ground-truth tokens at every position. The model never receives its own output as input during training. The hidden state is re-initialized at the start of each batch; for this training setup we treat each window independently rather than carrying state across windows.

The gradient clipping call clip_grad_norm_(model.parameters(), max_norm=5.0) is a practical necessity for RNN training. Before the gradient update, the gradient vector is scaled down if its norm exceeds the max_norm threshold. This prevents a single large gradient from causing an unstable weight update. The value 5.0 is a common default; for LSTM and GRU models, which have more stable gradients than vanilla RNNs, this threshold is often raised to 1.0 or even omitted.

Generating TextLink Copied

After training, we can generate text by running the RNN in autoregressive mode: each predicted token becomes the input for the next step.

In[15]:

Code

def generate_text(model, seed_char, max_len=100, temperature=0.8, device="cpu"):
    """
    Generate text autoregressively from a seed character.
    temperature > 1 gives more random output; temperature < 1 gives more deterministic output.
    """
    model.train(False)
    chars_generated = [seed_char]

    # Encode seed and initialize hidden state
    x = torch.tensor(
        [[char_to_idx[seed_char]]], dtype=torch.long, device=device
    )
    h = model.init_hidden(1, device)

    with torch.no_grad():
        for _ in range(max_len - 1):
            logits, h = model(x, h)
            # logits shape: (1, 1, vocab_size) -> squeeze to (vocab_size,)
            logits_squeezed = logits.squeeze() / temperature
            probs = torch.softmax(logits_squeezed, dim=-1)
            next_idx = torch.multinomial(probs, num_samples=1).item()
            next_char = idx_to_char[next_idx]
            chars_generated.append(next_char)
            x = torch.tensor([[next_idx]], dtype=torch.long, device=device)

    return "".join(chars_generated)


# Generate from two different seeds
generated_t = generate_text(model, "t", max_len=120, temperature=0.7)
generated_r = generate_text(model, "r", max_len=120, temperature=0.7)

Out[16]:

Console

Generated text (seed='t'):
  the quick brown fox jumps over the lazy dog recurrent neural networks process sequences step by step the hidden state ca

Generated text (seed='r'):
  rown fox jumps over the lazy dog recurrent neural networks process sequences step by step the hidden state carries conte

The temperature parameter controls how peaked the output distribution is. At temperature 1.0, the raw softmax probabilities are used. At temperature below 1.0, the distribution is sharpened (more deterministic), favoring the most probable tokens. At temperature above 1.0, the distribution is flattened (more random), which can produce more creative but less coherent text.

Notice how the generation loop passes h from one step to the next. This is the key difference between generation and training: at inference time, the hidden state must carry context across steps because there is no ground truth to reset to. The model builds up a context vector in h as it generates each character, and that accumulated context influences every subsequent character. If the hidden state forgets early parts of the generation (the vanishing gradient problem), the generated text may lose coherence over long spans.

Key Parameters of nn.RNNLink Copied

The key parameters for nn.RNN are:

input_size: Dimensionality of each input token's embedding vector. Larger values can encode richer token representations.
hidden_size: Number of hidden units $H$ . This is the primary capacity control, but note the quadratic parameter cost from $\mathbf{W}_{hh}$ .
num_layers: Number of stacked RNN layers. With num_layers=2, the hidden state of the first layer is fed as input to the second layer at each time step.
nonlinearity: Either 'tanh' (default) or 'relu'. The tanh option is standard for the reasons discussed above.
batch_first: Whether input tensors are (batch, seq, feature) or (seq, batch, feature). Setting batch_first=True is more intuitive for most training loops.
dropout: Dropout probability applied between RNN layers (not within a layer's recurrent connection). Only effective when num_layers > 1.
bidirectional: Whether to run a second RNN over the sequence in reverse. The output dimension doubles when set to True.