🎓 A1Learner
🔍

PyTorch GRU Cells

A GRU is a recurrent network that guards its hidden state with two gates instead of three. It keeps one state where an LSTM keeps two.

Gate rows side by side

The input weight matrix grows with the gate count. The plain RNN has one computation, the GRU three, and the LSTM four.

for name, cell in (("RNN", nn.RNN), ("GRU", nn.GRU), ("LSTM", nn.LSTM)):
    m = cell(input_size=1, hidden_size=16, batch_first=True)
    print(name, tuple(m.weight_ih_l0.shape))

One state comes back

A GRU call returns the output and the hidden state. No cell state rides along, unlike the LSTM pair.

output, hidden = gru(x)

Example

import torch
from torch import nn


def make_windows(wave, width):
    inputs = []
    targets = []
    for start in range(len(wave) - width):
        inputs.append(wave[start:start + width])
        targets.append(wave[start + width])
    return torch.stack(inputs).unsqueeze(-1), torch.stack(targets).unsqueeze(-1)


class SineModel(nn.Module):
    def __init__(self, cell, hidden):
        super().__init__()
        self.cell = cell(1, hidden, batch_first=True)
        self.head = nn.Linear(hidden, 1)

    def forward(self, x):
        output, _ = self.cell(x)
        return self.head(output[:, -1])


def train_and_score(cell, x_train, y_train, x_test, y_test):
    torch.manual_seed(0)
    model = SineModel(cell, 16)
    loss_fn = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    for epoch in range(50):
        loss = loss_fn(model(x_train), y_train)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    with torch.no_grad():
        return loss_fn(model(x_test), y_test).item()


wave = torch.sin(torch.arange(300) * 0.1)
x, y = make_windows(wave, 12)
x_train, y_train = x[:180], y[:180]
x_test, y_test = x[180:], y[180:]

print("Model  Test MSE")
for name, cell in (("RNN", nn.RNN), ("LSTM", nn.LSTM), ("GRU", nn.GRU)):
    error = train_and_score(cell, x_train, y_train, x_test, y_test)
    print(name, " ", round(error, 4))

Key points

  • Reset gate: decides how much of the old state shapes the new candidate
  • Update gate: decides how much of the old state carries over unchanged
  • A GRU has no separate cell state, only the hidden state
  • Gate rows at hidden size 16: RNN 16, GRU 48, LSTM 64
  • Use nn.GRU as a drop-in cell with the same call shape as nn.RNN and nn.LSTM
  • Fewer gates mean fewer weights and often faster steps
  • On short tasks the three cells often land close, as the closing table shows

In this example, the training function from the last lesson trains all three cells on the same sine windows. Seed 0 resets before each build. The table prints one test MSE per cell. The RNN scores 0.0019 and the LSTM scores 0.002, matching the last lesson exactly. The GRU scores 0.0016, the lowest of the three here. All three sit close together on this short task. The honest rule is to try the GRU when a task feels LSTM-sized and time matters. Long stretches with material that must survive untouched still favour the LSTM and its separate cell state.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

How many gates does a GRU use to guard its state?