GRU Architecture › Comparing GRU and LSTM Performance

import torch
import torch.nn as nn


class GRUClassifier(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        self.hidden_size = hidden_size
        self.gru = nn.GRU(input_size, hidden_size, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        output, hidden = self.gru(x)
        return self.fc(hidden.squeeze(0))


class LSTMClassifier(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        self.hidden_size = hidden_size
        self.lstm = nn.LSTM(input_size, hidden_size, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        output, (hidden, cell) = self.lstm(x)
        return self.fc(hidden.squeeze(0))


def generate_memory_task_data(
    batch_size, seq_length, signal_position, num_classes=8
):
    """Generate data for a memory task."""
    x = torch.randn(batch_size, seq_length, num_classes) * 0.1
    labels = torch.randint(0, num_classes, (batch_size,))
    for i in range(batch_size):
        x[i, signal_position, :] = 0
        x[i, signal_position, labels[i]] = 1.0
    return x, labels


def train_and_evaluate(model, seq_length, signal_position, epochs=100):
    """Train model and return final accuracy."""
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    criterion = nn.CrossEntropyLoss()

    for epoch in range(epochs):
        model.train()
        x, labels = generate_memory_task_data(32, seq_length, signal_position)
        optimizer.zero_grad()
        outputs = model(x)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

    model.eval()
    with torch.no_grad():
        x, labels = generate_memory_task_data(100, seq_length, signal_position)
        outputs = model(x)
        predictions = outputs.argmax(dim=1)
        accuracy = (predictions == labels).float().mean().item()

    return accuracy

# Compare GRU and LSTM on memory task
distances = [5, 10, 20, 30, 50, 75, 100]
seq_length = 110

torch.manual_seed(42)

gru_accuracies = []
lstm_accuracies = []

for distance in distances:
    signal_position = seq_length - distance - 1

    # Train GRU
    gru_model = GRUClassifier(input_size=8, hidden_size=64, output_size=8)
    gru_acc = train_and_evaluate(gru_model, seq_length, signal_position)
    gru_accuracies.append(gru_acc)

    # Train LSTM
    lstm_model = LSTMClassifier(input_size=8, hidden_size=64, output_size=8)
    lstm_acc = train_and_evaluate(lstm_model, seq_length, signal_position)
    lstm_accuracies.append(lstm_acc)