Building a Stock Price Forecasting Model in PyTorch with LSTMs

In this article, we will dive deep into how to build a stock price forecasting model using PyTorch and LSTM (Long Short-Term Memory) networks. LSTMs are a type of recurrent neural network (RNN) that are particularly effective for time series predictions due to their ability to capture long-term dependencies in sequential data.

What is an LSTM Network?

LSTM networks are a special kind of RNN, capable of learning long-term dependencies. They were introduced to solve the vanishing gradient problem typically encountered by traditional RNNs. LSTMs can maintain a "memory" over long sequences, enabling them to learn patterns from time series data over tens or hundreds of days.

Setting Up the Environment

Before we start, ensure that you have the necessary libraries installed. You will need PyTorch, pandas, numpy, and matplotlib. You can install these libraries using pip:

pip install torch pandas numpy matplotlib

Preparing the Data

We'll start by acquiring historical stock price data. You can use APIs from financial services like Alpha Vantage or Yahoo Finance to download data. For simplicity, we'll assume you already have a CSV file with the date, open, high, low, close, and volume data of the stock.

import pandas as pd

data = pd.read_csv('stock_prices.csv')
data['Date'] = pd.to_datetime(data['Date'])
data.sort_values('Date', inplace=True)

Next, we normalize the close prices, which is a crucial step for training neural networks as it often leads to quicker convergence:

import numpy as np

# Normalizing the data
close_prices = data['Close'].values
normalized_data = (close_prices - np.min(close_prices)) / (np.max(close_prices) - np.min(close_prices))

Creating the Training and Test Datasets

Time series data is typically divided into training and test datasets. Often, data until a specific date is used for training, and data beyond that date is used for testing the model's performance.

train_size = int(len(normalized_data) * 0.8)
train_data = normalized_data[0:train_size]
test_data = normalized_data[train_size:len(normalized_data)]

We'll now prepare the data for the LSTM model, which requires data to be structured in a particular way, usually in sequences:

def create_sequences(data, seq_length):
    sequences = []
    for i in range(len(data) - seq_length):
        seq = data[i:i + seq_length]
        label = data[i + seq_length]
        sequences.append((seq, label))
    return sequences

seq_length = 5  # Number of days to base prediction on
train_sequences = create_sequences(train_data, seq_length)
test_sequences = create_sequences(test_data, seq_length)

Building the LSTM Model

Now, let’s build our LSTM model in PyTorch. We'll define a class and initialize the layers:

import torch
import torch.nn as nn

class LSTMNetwork(nn.Module):
    def __init__(self):
        super(LSTMNetwork, self).__init__()
        self.lstm = nn.LSTM(input_size=1, hidden_size=50, num_layers=2, batch_first=True)
        self.fc = nn.Linear(50, 1)

    def forward(self, x):
        lstm_out, _ = self.lstm(x)
        predictions = self.fc(lstm_out[:, -1])
        return predictions

Training the Model

Before we start training, we need to convert our data into PyTorch tensors:

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = LSTMNetwork().to(device)

train_features = torch.tensor([s[0] for s in train_sequences], dtype=torch.float32).to(device)
train_labels = torch.tensor([s[1] for s in train_sequences], dtype=torch.float32).unsqueeze(1).to(device)

We'll use Mean Squared Error as our loss function and the Adam optimizer:

criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

Finally, let's train the model:

epochs = 100
def train(model, train_features, train_labels):
    model.train()
    for epoch in range(epochs):
        optimizer.zero_grad()
        predictions = model(train_features)
        loss = criterion(predictions, train_labels)
        loss.backward()
        optimizer.step()
        if epoch % 10 == 0:
            print(f'Epoch {epoch}, Loss: {loss.item()}')
train(model, train_features, train_labels)

Evaluating the Model

After training, it's crucial to evaluate how well the model predicts unseen data.

model.eval()
with torch.no_grad():
    test_features = torch.tensor([s[0] for s in test_sequences], dtype=torch.float32).to(device)
    test_predictions = model(test_features).cpu().numpy()

Finally, plot the results:

import matplotlib.pyplot as plt

dates = data['Date'].values
plt.plot(dates[train_size:len(normalized_data)], test_data, label='True Values')
plt.plot(dates[train_size:len(normalized_data)], test_predictions, label='Predicted Values')
plt.xlabel('Date')
plt.ylabel('Normalized Close Price')
plt.title('Stock Price Forecasting')
plt.legend()
plt.show()

Conclusion

Building a stock price forecasting model with LSTMs in PyTorch can be a robust way to predict future stock performance. Such models should be used with care, recognizing the stochastic nature of stock prices and the possibility of incorporating more sophisticated features to improve prediction accuracy. Excelling in this area can contribute valuable insights into market trends and investment strategies.