An open-source repository where I implemented RNN, LSTM, and GRU cells from scratch in PyTorch, then built a sentiment classifier and a language model on top of them. Nothing here delegates to PyTorch’s built-in recurrent modules: every gate is an explicit linear layer and every timestep is an explicit loop, so the forward pass can be read straight through. I wrote it as learning material, and that decision drove one tradeoff that runs through the whole codebase, which is that wherever clarity and speed conflict, clarity wins.
Three recurrent cells written by hand
The vanilla RNN cell concatenates the incoming hidden state with the input, pushes the result through a single linear layer and a ReLU to produce the next hidden state, then through a second linear layer for the output. The LSTM cell carries a separate cell state alongside the hidden state and guards it with four gates, forget, input, cell, and output, and the reason it handles long sequences is that the cell state is only ever touched by element-wise multiplication and addition rather than by a linear layer, so gradients flowing back through it are not repeatedly squashed. The GRU cell drops the separate cell state entirely and gets by with two gates, relevance and update. It blends old and new hidden state directly, in proportion to what the update gate decides.
Each cell is wrapped in a module that owns the timestep loop, initializes a zero hidden state on the input’s own device when none is supplied, and stacks the per-timestep outputs into a single tensor, with input shaped as batch by sequence length by features throughout.
Choosing clarity over speed in the LSTM
The LSTM cell allocates four separate linear layers, one per gate, and calls them one at a time. Every production implementation instead fuses those four into a single matrix multiply of four times the hidden size and slices the result, which is materially faster on a GPU. I left a comment directly above the gate definitions saying exactly that: this version is the most intuitive rather than the most efficient, and the four multiplications could be condensed into one.
That tradeoff belongs in the code rather than buried somewhere else. Someone reading forget_gate, input_gate, cell_gate, and output_gate as four named attributes can map them line by line onto the equations in the original paper, which is not possible on a first pass through a fused matmul and four tensor slices. The same choice shows up in the timestep loops, which concatenate outputs on every iteration instead of preallocating, slower and considerably easier to follow.
A callback training loop in pure PyTorch
The training loop reimplements fast.ai’s learner pattern with no fast.ai dependency. The learner holds the model, both dataloaders, the loss function, the learning rate, the optimizer factory, and a list of callbacks, and exposes six hook points fired around fit, epoch, and batch. The entire dispatch mechanism is a loop over the callbacks that fetches the named method and falls back to a no-op when it is missing. Callbacks therefore implement only the hooks they care about, with no base class to inherit from and no abstract methods to stub out.
The batch step runs the forward pass and computes the loss unconditionally but only zeroes gradients, backpropagates, and steps the optimizer when the model is in training mode, so the same code path serves training and validation. The epoch step takes a flag that selects the dataloader and toggles train or eval accordingly, and fit runs a training epoch followed by a validation epoch under no-grad. Three callbacks ship with it, one moving batches to the GPU and two tracking results, one for each task, with logging through Python’s standard logging module and optional Weights and Biases tracking for loss and metric curves.
Swapping architectures from a single config string
Two dispatch tables key off one string in the config file. The first maps names like lstmsentimentclf or rnnlanguagemodel to constructed models, the second maps those same names to the right callback list, so changing one line swaps the architecture and its metric tracking together. The config also carries batch size, epoch count, dataset, loss, learning rate, and optimizer, which means a run is fully described by a single file and there are no command-line arguments to remember.
The models are built to be swapped as well. The sentiment classifier takes the recurrent type as a plain string and resolves it against the module namespace at construction, so the same classifier runs on the hand-written RNN, LSTM, or GRU without a line of its own code changing.
Two tasks built on the cells
Sentiment classification trains on the IMDb reviews dataset. An embedding layer encodes tokens, the recurrent module produces an output per timestep, and dropout at 0.2 is applied across them. Rather than taking only the final output vector, which throws away most of the sequence, it takes both the average and the maximum across timesteps and concatenates them, which is why the linear head accepts twice the hidden size. No sigmoid is applied at the end, because the loss function folds it in for numerical stability.
Language modeling trains on fast.ai’s human-numbers benchmark, the first ten thousand integers written out as English words. It is deliberately tiny and highly regular, which makes it a useful sanity check on a hand-written cell, since a model that cannot learn to count places the bug in the cell rather than in the data. An embedding feeds the recurrent module and a linear layer projects each output to vocabulary size for the next-token distribution.
Training writes the model, the tokenizer, and the logs into a run directory, and a second entry point reads the same config to load those artifacts back for interacting with a trained model. Defaults are batch size 128, five epochs, Adam at 1e-3. Released under an MIT license.
Related reading
Companion deep-dives on two of these architectures, covering the theory that the code implements:
For a from-scratch implementation of the architecture that replaced these, see Transformer paper from scratch.
View the full code on GitHub →