Building LSTMs with PyTorch and Lightning AI Part 1: First Steps with LSTMs

In this article, we will explore how to implement an LSTM using PyTorch and Lightning.

For more details about LSTMs, there is a separate series of articles available here.

To begin, we first import the required modules.

import torch
import torch.nn as nn
import torch.nn.functional as F

We also introduce a new optimizer:

from torch.optim import Adam

Adam is used to fit the neural network to the data.

It works similarly to SGD, but in practice, Adam often converges faster and adapts the learning rate more effectively.

Next, we continue with the remaining imports:

import lightning as L
from torch.utils.data import TensorDataset, Data

We define the neural network by creating a Lightning module.

class LSTMByHand(L.LightningModule):
    def __init__(self):

    def lstm_unit(self, input_value, long_memory, short_memory):

    def forward(self, input):

    def configure_optimizers(self):

    def training_step(self, batch, batch_idx):

Now let’s implement the __init__

method.

This is where we initialize all weights and biases.

class LSTMByHand(L.LightningModule):
    def __init__(self):
        super().__init__()

        mean = torch.tensor(0.0)  # Mean of the normal distribution
        std = torch.tensor(1.0)   # Standard deviation

        self.wlr1 = nn.Parameter(torch.normal(mean=mean, std=std), requires_grad=True)
        self.wlr2 = nn.Parameter(torch.normal(mean=mean, std=std), requires_grad=True)
        self.blr1 = nn.Parameter(torch.tensor(0.0), requires_grad=True)

        self.wpr1 = nn.Parameter(torch.normal(mean=mean, std=std), requires_grad=True)
        self.wpr2 = nn.Parameter(torch.normal(mean=mean, std=std), requires_grad=True)
        self.bpr1 = nn.Parameter(torch.tensor(0.0), requires_grad=True)

        self.wp1 = nn.Parameter(torch.normal(mean=mean, std=std), requires_grad=True)
        self.wp2 = nn.Parameter(torch.normal(mean=mean, std=std), requires_grad=True)
        self.bp1 = nn.Parameter(torch.tensor(0.0), requires_grad=True)

        self.wo1 = nn.Parameter(torch.normal(mean=mean, std=std), requires_grad=True)
        self.wo2 = nn.Parameter(torch.normal(mean=mean, std=std), requires_grad=True)
        self.bo1 = nn.Parameter(torch.tensor(0.0), requires_grad=True)

Unlike earlier examples, we initialize weights using a normal distribution.

Before moving further, let’s understand what that means.

Imagine measuring the heights of a large group of people:

When plotted, this forms a symmetric bell-shaped curve.

This is called a normal distribution.

We use:

0

1

Also, all parameters have requires_grad=True

, meaning they will be trained during backpropagation.

Next, we will explore the lstm_unit

function and how the LSTM actually processes information step by step.

AI agents write code fast. They also silently remove logic, change behavior, and introduce bugs -- without telling you. You often find out in production.

git-lrc 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