Teaching PyTorch to predict fairy tales, one character at a time

Char-rnn, the character-level recurrent network popularized by Andrej Karpathy, is a great first project for anyone who wants to see a neural network do something tangible. Given enough text, the model learns which character likely comes next, and after enough training it starts producing output that resembles the source material.

In a recent experiment, I trained a PyTorch LSTM on Hans Christian Andersen's fairy tales from Project Gutenberg. The results weren't as polished as Karpathy's demos, but the process of building and debugging the model was illuminating. Below is a walkthrough of the approach, the code, and the output. The full Jupyter notebook is available as a gist; with Colab you can run it on a free GPU.

Preparing the data

The training data is a plain-text file of fairy tales, downloaded from Project Gutenberg. The key preprocessing step converts characters into numbers that a network can consume. For this, I used the Vocab class from fastai, which builds a character-to-integer mapping.

!wget -O fairy-tales.txt

After running the vocabulary over the text, the entire corpus becomes a single large array of integer tokens called training_set. This numeric representation is what we feed into the network during training.

from fastai.text import *
text = unidecode.unidecode(open('fairy-tales.txt').read())
v = Vocab.create((x for x in text), max_vocab=400, min_freq=1)
training_set = torch.Tensor(v.numericalize([x for x in text])).type(torch.LongTensor).cuda()
num_letters = len(v.itos)

Building a character-level LSTM

The model is a thin wrapper around PyTorch's built-in nn.LSTM class, with a few adjustments to make it work for character prediction:

  1. Input characters are one-hot encoded so their dimension matches the LSTM's expected input size.
  2. A linear layer following the LSTM maps the hidden state (size hidden_size) back to a vector of size input_size—one score per possible character.
  3. The hidden and cell states are stored as instance variables and detached after every batch; calling .detach() stops gradient tracking at that point, which is necessary to avoid backpropagating through the entire training history.
class MyLSTM(nn.Module):
    def __init__(self, input_size, hidden_size):
        super().__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, batch_first=True)
        self.h2o = nn.Linear(hidden_size, input_size)
        self.input_size=input_size
        self.hidden = None
        
    def forward(self, input):
        input = torch.nn.functional.one_hot(input, num_classes=self.input_size).type(torch.FloatTensor).cuda().unsqueeze(0)
        if self.hidden is None:
            l_output, self.hidden = self.lstm(input)
        else:
            l_output, self.hidden = self.lstm(input, self.hidden)
        self.hidden = (self.hidden[0].detach(), self.hidden[1].detach())
        
        return self.h2o(l_output)

One subtlety worth noting: if you pass a sequence of tokens (say, six letters) into nn.LSTM at once, it internally rolls the hidden state forward for each token and computes gradients across the full sequence using backpropagation through time (BPTT). You don't need to write the loop yourself.

Training loop

My first attempt used fastai's training helpers, but the abstraction made it hard to trace what was happening. Writing the training loop by hand was clearer. Each iteration does the following:

  1. Feed the network a character sequence such as and they ought not to teas, encoded as integers.
  2. Take the model's prediction for the character that should follow.
  3. Compare that prediction against the actual next character (e in the example above) using a loss function.
  4. Call loss.backward() to compute gradients.
  5. Update the weights with self.optimizer.step().
class Trainer():
  def __init__(self):
      self.rnn = MyLSTM(input_size, hidden_size).cuda()
      self.optimizer = torch.optim.Adam(self.rnn.parameters(), amsgrad=True, lr=lr)
  def epoch(self):
      i = 0
      while i < len(training_set) - 40:
        seq_len = random.randint(10, 40)
        input, target = training_set[i:i+seq_len],training_set[i+1:i+1+seq_len]
        i += seq_len
        # forward pass
        output = self.rnn(input)
        loss = F.cross_entropy(output.squeeze()[-1:], target[-1:])
        # compute gradients and take optimizer step
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

My initial implementation attempted manual BPTT: I passed one character at a time and only took an optimization step after 20 characters. While the loss initially decreased, it later spiked unpredictably. Switching to the simpler approach—feeding all 20 characters at once as the seq_len dimension and letting PyTorch's LSTM handle the backpropagation internally—produced much more stable training.

for i in range(20):
    input, target = next(iter)
    output, hidden = self.lstm(input, hidden)
loss = F.cross_entropy(output, target)
hidden = hidden.detach()
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()

After roughly 300 passes over the same data (about an hour of training) the model started producing recognizable English-like sequences. I did not bother with validation set split, since overfitting wasn't a concern for this experiment.

Generating text

Generation loops over the model's predictions. The helper functions make_preds and next_pred handle the bookkeeping; the core operation is:

output = rnn(input)
prediction_vector = F.softmax(output/temperature)
letter = v.textify(torch.multinomial(prediction_vector, 1).flatten(), sep='').replace('_', ' ')

Each step works like this:

  1. The LSTM outputs an unnormalized score vector, one entry per possible character.
  2. F.softmax(output / temperature) converts those scores into probabilities. The temperature parameter controls how peaked the distribution is—a tiny value (like 0.0000001) makes the model always choose the highest-probability character.
  3. torch.multinomial(prediction_vector) samples an index according to that probability distribution.
  4. v.textify maps the sampled index back to a character.

Repeat this process a few hundred times, and you have a full generated passage.

Sampling the results

With temperature = 1, the output has the rhythm and spelling of English—punctuation, capitalization, and sentence breaks land in plausible places—but the content is pure nonsense:

“An who was you colotal said that have to have been a little crimantable and beamed home the beetle. “I shall be in the head of the green for the sound of the wood. The pastor. “I child hand through the emperor’s sorthes, where the mother was a great deal down the conscious, which are all the gleam of the wood they saw the last great of the emperor’s forments, the house of a large gone there was nothing of the wonded the sound of which she saw in the converse of the beetle. “I shall know happy to him. This stories herself and the sound of the young mons feathery in the green safe.”
“That was the pastor. The some and hand on the water sound of the beauty be and home to have been consider and tree and the face. The some to the froghesses and stringing to the sea, and the yellow was too intention, he was not a warm to the pastor. The pastor which are the faten to go and the world from the bell, why really the laborer’s back of most handsome that she was a caperven and the confectioned and thoughts were seated to have great made

Turning the temperature down to 0.1 makes the distribution much more confident, and the output repeats familiar phrases over and over:

ole the sound of the beauty of the beetle. “She was a great emperor of the sea, and the sun was so warm to the confectioned the beetle. “I shall be so many for the beetle. “I shall be so many for the beetle. “I shall be so standen for the world, and the sun was so warm to the sea, and the sun was so warm to the sea, and the sound of the world from the bell, where the beetle was the sea, and the sound of the world from the bell, where the beetle was the sea, and the sound of the wood flowers and the sound of the wood, and the sound of the world from the bell, where the world from the wood, and the sound of the

Thematically, the model is fixated on beetles, confectioners, the sea, and the sun. It's a strange obsession, but given that it learned English purely from character sequences, the output is still impressive.

Where the model underperforms

The output is far less coherent than what Karpathy demonstrated. Possible reasons include:

  • Too little training data available.
  • Training was cut short—about an hour because I grew impatient with babysitting the Colab session.
  • Karpathy's network used two LSTM layers with wider hidden states; mine was a single layer.
  • The instability might have another cause I haven't diagnosed.

Despite those caveats, the model learned enough to generate vaguely grammatical text from scratch—a solid outcome for a first neural network project.