🎓 A1Learner
🔍

PyTorch Sequence Batching

Sequence batching stacks several sequences into one tensor with three axes. The axes are batch, time, and features, in that order when batch_first is set.

Time first by default

A plain nn.RNN expects time before batch. Three sequences of length 10 give output shape (10, 3, 4).

plain = nn.RNN(input_size=1, hidden_size=4)
x = torch.zeros(10, 3, 1)
out, h = plain(x)
print(tuple(x.shape), tuple(out.shape))

Move the batch first

Swapping the first two axes turns the same data into batch-first order.

print(tuple(x.permute(1, 0, 2).shape))

Output holds every step

The output keeps one row per step for the whole batch. The hidden tensor keeps only the final state.

output, hidden = rnn(batch)
print(tuple(output.shape), tuple(hidden.shape))

Example

import torch
from torch import nn

torch.manual_seed(0)
wave = torch.sin(torch.arange(200) * 0.1)
width = 10
windows = torch.stack([wave[s:s + width] for s in (0, 40, 80)])
batch = windows.unsqueeze(-1)

rnn = nn.RNN(input_size=1, hidden_size=4, batch_first=True)
with torch.no_grad():
    output, hidden = rnn(batch)

print("Input shape (batch, time, features):", tuple(batch.shape))
print("Output shape:", tuple(output.shape))
print("Hidden shape (layers, batch, hidden):", tuple(hidden.shape))
print("Last-step output row 0:", [round(v, 3) for v in output[0, -1].tolist()])
print("Hidden row 0:", [round(v, 3) for v in hidden[0, 0].tolist()])
print("They match:", torch.allclose(output[:, -1], hidden[0]))

Key points

  • batch_first=True puts the batch axis first in input and output
  • Input shape: (batch, time, features) under batch_first
  • Output shape: (batch, time, hidden) with one row per step
  • Hidden shape: (layers, batch, hidden), and it never follows batch_first
  • The output at the last step equals the final hidden state for a plain RNN
  • Without batch_first, the input axes are (time, batch, features)

In this example, three sine windows of width 10 form the batch. The unsqueeze adds a feature axis of size 1, so the input shape is (3, 10, 1). The RNN reads the batch with hidden size 4. The output shape is (3, 10, 4), one hidden row for each of the ten steps. The hidden shape is (1, 3, 4), one layer, three sequences, four numbers. The last-step output row and the hidden row print the same four numbers. The closing check confirms they match. A model that predicts from a whole sequence reads that final row. It can read the hidden state instead and ignore the rest of the output.

Note

The hidden tensor keeps its layers-first order in both conventions

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Under batch_first, an input of shape (3, 10, 1) means what?