Encoder-Decoder Framework › PyTorch Seq2Seq Implementation
Language AI HandbookAll chapters Encoder-Decoder Framework › PyTorch Seq2Seq ImplementationOne section · Sections in this chapter10 / 16 ↓
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
Let us now build a complete seq2seq model from scratch using PyTorch. We will implement it for a character-level sequence transformation task (mapping phrases to their word-reversed versions) to keep the vocabulary small and training fast.
Setup and Data Preparation
We start by importing all required libraries and preparing training data.
import random
import numpy as np
import torch
# Set seeds for reproducibility
random.seed(42)
np.random.seed(42)
torch.manual_seed(42)
# Special token indices
PAD_IDX = 0
SOS_IDX = 1
EOS_IDX = 2
# Build character vocabulary (lowercase letters + special tokens)
chars = ["<PAD>", "<SOS>", "<EOS>"] + list("abcdefghijklmnopqrstuvwxyz ")
char_to_idx = {ch: i for i, ch in enumerate(chars)}
idx_to_char = {i: ch for ch, i in char_to_idx.items()}
VOCAB_SIZE = len(chars)def make_sample(phrase):
"""Source: original phrase. Target: words in reversed order."""
words = phrase.split()
reversed_phrase = " ".join(reversed(words))
return phrase, reversed_phrase
# Training sentences
train_phrases = [
"the cat sat on the mat",
"a dog ran in the park",
"she reads books every day",
"we learn new things today",
"the bird sang a song",
"he writes code at night",
"they swim in the lake",
"i drink tea in the morning",
"the sun shines very bright",
"rain falls on green leaves",
"good food makes people happy",
"fast cars drive on roads",
"children play in the yard",
"music sounds so very nice",
"flowers bloom in warm spring",
"stars shine in dark skies",
"rivers flow to the sea",
"trees grow in deep forests",
"snow falls in cold winter",
"bread bakes in warm ovens",
]
train_pairs = [make_sample(p) for p in train_phrases]def encode_sequence(text, char_to_idx, max_len=35):
"""Encode a text string to a padded tensor with SOS and EOS."""
indices = (
[SOS_IDX] + [char_to_idx.get(ch, PAD_IDX) for ch in text] + [EOS_IDX]
)
indices = indices[:max_len]
indices += [PAD_IDX] * (max_len - len(indices))
return torch.tensor(indices, dtype=torch.long)
MAX_LEN = 35
src_tensors = [
encode_sequence(src, char_to_idx, MAX_LEN) for src, tgt in train_pairs
]
tgt_tensors = [
encode_sequence(tgt, char_to_idx, MAX_LEN) for src, tgt in train_pairs
]
src_batch = torch.stack(src_tensors) # (N, max_len)
tgt_batch = torch.stack(tgt_tensors) # (N, max_len)Vocabulary size: 30 Training pairs: 20 Source batch shape: torch.Size([20, 35]) Target batch shape: torch.Size([20, 35]) Sample pair: Source: 'the cat sat on the mat' Target: 'mat the on sat cat the' Encoded source (first 10): [1, 22, 10, 7, 29, 5, 3, 22, 29, 21] Decoded source: 'the cat sat on the mat'
The training data maps English phrases to their word-reversed versions, which is a minimal but illustrative sequence transformation task. With 29 vocabulary items (26 letters plus space and special tokens), the model can focus on learning the seq2seq mechanics rather than a large vocabulary. Note that the character-level encoding means the model treats each individual character as a token, not each word, so "the cat sat on the mat" becomes a sequence of 22 character tokens plus the start and end tokens.
Encoder
The encoder wraps an embedding layer and an LSTM. It reads all source tokens and returns the final hidden and cell states.
import torch.nn as nn
class Encoder(nn.Module):
def __init__(
self, vocab_size, embed_dim, hidden_dim, n_layers=1, dropout=0.3
):
super().__init__()
self.embedding = nn.Embedding(
vocab_size, embed_dim, padding_idx=PAD_IDX
)
self.rnn = nn.LSTM(
embed_dim,
hidden_dim,
n_layers,
batch_first=True,
dropout=dropout if n_layers > 1 else 0,
)
self.dropout = nn.Dropout(dropout)
def forward(self, src):
# src: (batch, seq_len)
embedded = self.dropout(
self.embedding(src)
) # (batch, seq_len, embed_dim)
outputs, (hidden, cell) = self.rnn(
embedded
) # outputs: (batch, seq_len, hidden_dim)
# Return only final hidden and cell states as context
return hidden, cellThe forward method returns the LSTM's final (hidden, cell) tuple, which becomes the context vector. We discard the intermediate hidden states outputs because the basic seq2seq model uses only the final state. The attention-based models in later chapters will use all intermediate states: the attention mechanism needs access to to form the context-sensitive weighted sum at each decoding step.
The padding_idx argument to nn.Embedding sets the embedding vector for the pad token to all zeros and keeps its gradient at zero during training. This ensures that padded positions do not affect the learned embeddings for real tokens.
Decoder
The decoder also wraps an embedding layer and an LSTM, but it generates one token at a time and uses the context vector as its initial hidden state.
class Decoder(nn.Module):
def __init__(
self, vocab_size, embed_dim, hidden_dim, n_layers=1, dropout=0.3
):
super().__init__()
self.embedding = nn.Embedding(
vocab_size, embed_dim, padding_idx=PAD_IDX
)
self.rnn = nn.LSTM(
embed_dim,
hidden_dim,
n_layers,
batch_first=True,
dropout=dropout if n_layers > 1 else 0,
)
self.fc_out = nn.Linear(hidden_dim, vocab_size)
self.dropout = nn.Dropout(dropout)
def forward(self, input_token, hidden, cell):
# input_token: (batch,)
input_token = input_token.unsqueeze(1) # (batch, 1)
embedded = self.dropout(
self.embedding(input_token)
) # (batch, 1, embed_dim)
output, (hidden, cell) = self.rnn(
embedded, (hidden, cell)
) # (batch, 1, hidden_dim)
prediction = self.fc_out(output.squeeze(1)) # (batch, vocab_size)
return prediction, hidden, cellThe decoder processes one token at a time. It takes a single token input_token, the current hidden state hidden, and cell state cell, and produces logits over the vocabulary plus updated states. This single-step design lets us implement teacher forcing easily by controlling which token we feed in at each step.
The fc_out layer is a linear projection from the hidden dimension to the full vocabulary size. At each step, the decoder's hidden state encodes both what has been generated so far and what the encoder told it about the source. The linear projection maps this rich representation to an unnormalized score for every possible next token, and the softmax during training (handled internally by nn.CrossEntropyLoss) converts these to probabilities.
Seq2Seq Model
The seq2seq wrapper combines the encoder and decoder and implements teacher forcing.
class Seq2Seq(nn.Module):
def __init__(self, encoder, decoder, device):
super().__init__()
self.encoder = encoder
self.decoder = decoder
self.device = device
def forward(self, src, tgt, teacher_forcing_ratio=0.5):
# src: (batch, src_len), tgt: (batch, tgt_len)
batch_size = src.shape[0]
tgt_len = tgt.shape[1]
# Store decoder outputs (logits over vocabulary at each step)
outputs = torch.zeros(batch_size, tgt_len, VOCAB_SIZE).to(self.device)
# Encode source sequence
hidden, cell = self.encoder(src)
# First decoder input is the SOS token
input_token = tgt[:, 0] # (batch,)
for t in range(1, tgt_len):
output, hidden, cell = self.decoder(input_token, hidden, cell)
outputs[:, t, :] = output
use_teacher_forcing = random.random() < teacher_forcing_ratio
if use_teacher_forcing:
input_token = tgt[:, t] # Ground-truth token
else:
input_token = output.argmax(dim=1) # Model's best guess
return outputsThe teacher_forcing_ratio controls the mix: at 1.0 we always use ground truth (pure teacher forcing), at 0.0 we always use the model's own predictions (fully autoregressive), and values in between give a blend. Notice that the loop starts at , not . Position 0 in the target is the <SOS> token, which is the initial decoder input but not a target label. The output at position is the model's prediction for target position , so we align prediction at with label at .
This forward pass accumulates the full output tensor at each decoding step, which is what we need for computing the loss during training. At inference time, we do not need to precompute the full output tensor; we can greedily pick the best token at each step and stop when <EOS> is generated.
Training Loop
import torch.optim as optim
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
EMBED_DIM = 64
HIDDEN_DIM = 128
N_LAYERS = 1
DROPOUT = 0.0
LEARNING_RATE = 5e-3
N_EPOCHS = 200
TEACHER_FORCING_RATIO = 1.0
encoder = Encoder(VOCAB_SIZE, EMBED_DIM, HIDDEN_DIM, N_LAYERS, DROPOUT).to(
device
)
decoder = Decoder(VOCAB_SIZE, EMBED_DIM, HIDDEN_DIM, N_LAYERS, DROPOUT).to(
device
)
model = Seq2Seq(encoder, decoder, device).to(device)
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
criterion = nn.CrossEntropyLoss(ignore_index=PAD_IDX)
src_batch_dev = src_batch.to(device)
tgt_batch_dev = tgt_batch.to(device)
train_losses = []
for epoch in range(N_EPOCHS):
model.train()
optimizer.zero_grad()
outputs = model(src_batch_dev, tgt_batch_dev, TEACHER_FORCING_RATIO)
# Ignore position 0 (SOS); align predictions with targets
out_flat = outputs[:, 1:, :].reshape(-1, VOCAB_SIZE)
tgt_flat = tgt_batch_dev[:, 1:].reshape(-1)
loss = criterion(out_flat, tgt_flat)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
train_losses.append(loss.item())Training complete over 200 epochs Initial loss: 3.4082 Final loss: 0.0035 Loss reduction: 3.4047
The loss drops substantially from its initial value toward near zero, confirming the model has learned to reproduce the training sequences. With only 20 training examples and teacher forcing enabled throughout, the model memorizes the training set quickly. The ignore_index=PAD_IDX argument to CrossEntropyLoss ensures that padded positions do not contribute to the gradient, while clip_grad_norm_ prevents gradient explosions that are common in LSTM training.
The reshape operations before computing the loss are important. The model outputs a tensor of shape (batch, tgt_len, vocab_size), but CrossEntropyLoss expects predictions of shape (N, vocab_size) and targets of shape (N,). Slicing off position 0 (which is the SOS token, not a target label) and then flattening aligns the prediction for step with the target label at position .
Inference
At inference time, there is no teacher forcing. The decoder uses its own predicted token at each step.
def translate(
model, source_text, char_to_idx, idx_to_char, max_len=35, device="cpu"
):
"""Translate a source string using greedy decoding."""
model.eval()
with torch.no_grad():
src_tensor = (
encode_sequence(source_text, char_to_idx, max_len)
.unsqueeze(0)
.to(device)
)
hidden, cell = model.encoder(src_tensor)
input_token = torch.tensor([SOS_IDX], device=device)
output_chars = []
for _ in range(max_len):
pred, hidden, cell = model.decoder(input_token, hidden, cell)
best_token = pred.argmax(dim=1)
char = idx_to_char[best_token.item()]
if best_token.item() == EOS_IDX:
break
if best_token.item() != PAD_IDX:
output_chars.append(char)
input_token = best_token
return "".join(output_chars)Inference results (greedy decoding): Source Expected Predicted ---------------------------------------------------------------------------------------------------- the cat sat on the mat mat the on sat cat the mat the on sat cat the a dog ran in the park park the in ran dog a park the in ran dog a she reads books every day day every books reads she day every books reads she
The model accurately reproduces the word-reversal transformation for the training sentences. This demonstrates that the encoder successfully compressed the source phrase into a context vector containing enough information for the decoder to produce the correct output. Notice that the inference function uses model.eval() and torch.no_grad(), which disables dropout and gradient computation respectively. In eval mode, dropout layers pass their inputs unchanged, and gradient tracking is disabled to save memory and computation.
Key Parameters
The key hyperparameters for the seq2seq model are:
- embed_dim: Dimensionality of token embeddings. Larger values capture richer representations but increase memory and computation.
- hidden_dim: Dimensionality of the LSTM hidden state. This is the size of the context vector, the fundamental bottleneck of the architecture.
- n_layers: Number of stacked LSTM layers. Multiple layers allow hierarchical representation, as covered in the Stacked RNNs chapter.
- teacher_forcing_ratio: Probability of using ground-truth tokens during training. Higher values speed convergence but increase exposure bias.
- dropout: Regularization for the embedding and LSTM layers. Helps prevent overfitting on larger datasets.