Building LSTMs with PyTorch and Lightning AI Part 8: Setting Up a Simpler LSTM

In the

Notice that all of the graphs have become much flatter, indicating that the model has nearly converged.

The predictions for Company A and Company B are now very close to their target values of 0 and 1, respectively.

At this point, we have successfully trained our LSTM model.

Now let's look at an even simpler way to build an LSTM using PyTorch's built-in nn.LSTM module.

We will start by creating another class.

class LightningLSTM(L.LightningModule):

    def __init__(self):

        super().__init__()

        self.lstm = nn.LSTM(input_size=1, hidden_size=1)

The input_size specifies the number of features in each input.

In our example, each day contains only a single feature, which is the stock price of a company.

The hidden_size specifies the number of values in the hidden state.

For this example, we use a hidden size of 1, since we ultimately want the model to predict a single value for Day 5.

Next, we will implement the forward() method, which works a little differently when using nn.LSTM. We will explore that in the next article.

fixes this. It hooks into git commit and reviews every diff before it lands. 60-second setup. Completely free.

Any feedback or contributors are welcome! It's online, source-available, and ready for anyone to use.

Give it a ⭐ star on Github