Skip to main content

PyTorch Implementation: LEGO Bricks

We've spent a lot of time talking about the intense math behind LSTMs and GRUs—VIP Highways, Sigmoid valves, Update gates, and Bidirectional reading.

You might be terrified that you have to write all of this math by hand.

Don't panic! The beauty of modern AI frameworks like PyTorch is that they package all of this complex math into simple, pre-built LEGO bricks. You don't need to know how to build the engine; you just need to know how to start the car.

1. Summoning an LSTM​

To create an LSTM in PyTorch, you only need one single line of code!

import torch
import torch.nn as nn

# Create an LSTM network!
# input_size: The size of your word vectors (e.g., 300 coordinates from Word2Vec)
# hidden_size: How big you want the AI's "memory clipboard" to be (e.g., 128)
my_lstm = nn.LSTM(input_size=300, hidden_size=128)

That's it! PyTorch automatically handles all the Forget, Input, and Output gates behind the scenes.

2. Passing Data Through​

When you want to pass a sentence (a sequence of word vectors) into your LSTM, it will return two things:

  1. The Outputs: The short-term memory for every single word in the sentence.
  2. The Final States: The very last VIP Highway (Cell State) and short-term memory (Hidden State) at the end of the sentence.
# Pass your sentence into the LSTM
# Remember, PyTorch expects a tensor (a grid of numbers)
outputs, (final_hidden, final_cell) = my_lstm(my_sentence_tensor)

3. Switching to a GRU​

What if you decide the LSTM is running too slowly, and you want to try the "2-in-1 Shampoo" GRU instead?

You literally just change four letters in your code.

# Create a GRU instead!
my_gru = nn.GRU(input_size=300, hidden_size=128)

# Pass the data through (Notice there is no final_cell, because GRUs don't have a VIP highway!)
outputs, final_hidden = my_gru(my_sentence_tensor)

4. Making it Bidirectional​

Do you want your AI to read the sentence both forwards and backwards so it understands context perfectly?

Just set bidirectional=True.

# Create a Bidirectional LSTM!
my_bilstm = nn.LSTM(input_size=300, hidden_size=128, bidirectional=True)

PyTorch automatically clones the network, runs one forward, runs one backward, and glues the outputs together for you. It feels almost like cheating!

Next Up: We've pushed Recurrent Neural Networks to their absolute limits. But what happens when we try to use an LSTM to translate a massive, 1,000-word essay? We hit a wall. Welcome to Chapter 5, where we break that wall by inventing Attention!