PyTorch - Recurrent Neural Network

PyTorch - Recurrent Neural Networks (RNNs)

Recurrent Neural Networks (RNNs) are a class of neural networks well-suited for processing sequential data. Unlike feedforward networks where inputs are assumed to be independent, RNNs maintain an internal state (or memory) that captures information about previous elements in the sequence. This makes them powerful for tasks like natural language processing, speech recognition, and time series analysis.

An RNN processes sequences by iterating through the sequence elements and updating its hidden state at each step. The computation at each time step t typically involves the current input x_t and the previous hidden state h_{t-1} to produce the new hidden state h_t and an optional output o_t. This recurrent loop allows information to persist across time steps.

In this chapter, we’ll demonstrate how to implement a simple RNN from scratch in PyTorch to model a sine wave. This will help in understanding the mechanics, though for practical applications, PyTorch provides optimized built-in RNN layers.

During training, we will feed the model one data point (time step) at a time. The input sequence x will consist of 20 data points, and the target sequence y will be the input sequence shifted by one time step, i.e., predicting the next point in the sequence.

Step 1: Imports and Setup

Import necessary packages. We’ll use torch for neural network functionalities, numpy for numerical operations, and matplotlib for plotting.

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
import torch.nn.init as init

Step 2: Hyperparameters and Data Preparation

Define model hyperparameters and generate the sine wave data. The input to our RNN at each time step will be the current data point. The hidden layer size determines the capacity of our RNN’s memory.

dtype = torch.float32
input_size = 1  # Input dimension (sine wave value at time t)
hidden_size = 6 # Size of the hidden state
output_size = 1 # Output dimension (predicted sine wave value at t+1)
epochs = 300
seq_length = 20 # Length of the input sequence
learning_rate = 0.1

# Generate sine wave data
data_time_steps = np.linspace(2, 10, seq_length + 1, dtype=np.float32)
data = np.sin(data_time_steps)
data.resize((seq_length + 1, 1))

# Create input and target sequences
x_np = data[:-1]
y_np = data[1:]

x = torch.tensor(x_np, dtype=dtype)
y = torch.tensor(y_np, dtype=dtype)

x is the input data sequence, and y is the target sequence (x shifted by one time step).

Step 3: Manual RNN Weight Initialization

We’ll manually define and initialize the weights for our simple RNN. w1 maps the concatenated input and previous hidden state to the new hidden state. w2 maps the new hidden state to the output prediction. We set requires_grad=True to enable gradient computation for these weights.

# Weights for input to hidden layer, and hidden to hidden layer (combined)
# Input x_t (size 1) and previous hidden h_{t-1} (size hidden_size) are concatenated
w1 = torch.empty(input_size + hidden_size, hidden_size, dtype=dtype)
init.xavier_normal_(w1) # A common initialization strategy
w1.requires_grad_(True)

# Weights for hidden to output layer
w2 = torch.empty(hidden_size, output_size, dtype=dtype)
init.xavier_normal_(w2)
w2.requires_grad_(True)

# Bias for the hidden layer (optional, can be added)
# b1 = torch.zeros(hidden_size, dtype=dtype, requires_grad=True)
# Bias for the output layer (optional)
# b2 = torch.zeros(output_size, dtype=dtype, requires_grad=True)

Step 4: Define the Forward Pass Function

This function defines one step of the RNN computation.

def forward_step(input_t, context_state, w1_param, w2_param):
    # Concatenate input_t and previous context_state (hidden_state)
    # Ensure input_t is 2D: (1, input_size)
    if input_t.ndim == 1:
        input_t = input_t.unsqueeze(0) # Make it (1, input_size)

    xh = torch.cat((input_t, context_state), 1)
    # New hidden state (context_state)
    context_state_next = torch.tanh(xh.mm(w1_param)) # Add + b1 if using bias
    # Output prediction
    out = context_state_next.mm(w2_param) # Add + b2 if using bias
    return out, context_state_next

Step 5: Training the RNN

The training loop iterates for a number of epochs. In each epoch, we process the entire sequence. We use Mean Squared Error (MSE) as the loss function.

for i in range(epochs):
    total_loss = 0
    # Initialize context_state (hidden state) for the beginning of the sequence
    current_context_state = torch.zeros((1, hidden_size), dtype=dtype)

    for j in range(x.size(0)): # Iterate through each time step in the sequence
        input_t = x[j:(j + 1)]   # Current input (shape: [1, input_size])
        target_t = y[j:(j + 1)] # Current target (shape: [1, output_size])

        pred_t, next_context_state = forward_step(input_t, current_context_state, w1, w2)

        loss = (pred_t - target_t).pow(2).sum() / 2 # MSE variant
        total_loss += loss.item()

        # Backward pass and update weights manually
        # First, clear old gradients for w1 and w2 if they exist from previous iterations
        if w1.grad is not None:
            w1.grad.zero_()
        if w2.grad is not None:
            w2.grad.zero_()

        loss.backward() # Compute gradients

        with torch.no_grad(): # Temporarily disable gradient tracking for updates
            w1 -= learning_rate * w1.grad
            w2 -= learning_rate * w2.grad
            # If using biases, update them here as well: b1 -= learning_rate * b1.grad

        # Detach the context_state to prevent gradients from flowing back to the beginning of time
        # This is crucial for Truncated Backpropagation Through Time (TBPTT) like behavior
        # or when processing sequence step-by-step manually.
        current_context_state = next_context_state.detach()

    if (i + 1) % 10 == 0:
        print(f"Epoch: {i+1}/{epochs}, Loss: {total_loss / x.size(0):.4f}")

# After training, generate predictions
predictions = []
with torch.no_grad(): # No need to track gradients for prediction
    current_context_state_pred = torch.zeros((1, hidden_size), dtype=dtype)
    for i in range(x.size(0)):
        input_val = x[i:(i+1)]
        pred, current_context_state_pred = forward_step(input_val, current_context_state_pred, w1, w2)
        predictions.append(pred.item())

Step 6: Plotting the Results

Plot the original sine wave and the RNN’s predictions.

plt.figure(figsize=(12, 6))
plt.title('Sine Wave Prediction with Manual RNN')
plt.xlabel('Time Step')
plt.ylabel('Value')
plt.plot(data_time_steps[:-1], x.numpy().flatten(), 'bo-', label='Actual Data (Input)')
plt.plot(data_time_steps[1:], np.array(predictions), 'ro--', label='Predicted Data')
plt.legend()
plt.grid(True)
plt.show()

Output Description

The output will be a plot showing the original sine wave (typically as blue circles or a line) and the RNN’s predicted sine wave (typically as red dashed line or circles). A good model will show the predicted wave closely following the actual wave.

Note: While this example demonstrates building an RNN from scratch for educational purposes, PyTorch provides optimized and convenient modules like torch.nn.RNN, torch.nn.LSTM, and torch.nn.GRU in its torch.nn package. These are recommended for most practical applications as they are efficient and handle many complexities internally.