I did a bunch of things to prep for implementing a wavenet architecture, which basically changes the model so instead of taking in all the characters of input at once and squishing them all into the first layer, it gradually puts them in at each separate layer. First, I put more stuff in classes! This didn’t really change the functionality, but instead of embedding the characters and flattening the model in separate chunks, I put them in their own classes so they can get put in the list of layers, and we can just run the input through each layer in the list, and compare it with the target at the end, instead of having half the layers in the list and half as separate chunks, which is a lot smoother. So the entire model is just:

model = Sequential([  
    Embedding(vocab_size, n_embd),
    FlattenConsecutive(2), Linear(n_embd * 2, n_hidden, bias=False), BatchNorm1D(n_hidden), Tanh(),  
    FlattenConsecutive(2), Linear(n_hidden * 2, n_hidden, bias=False), BatchNorm1D(n_hidden), Tanh(),  
    FlattenConsecutive(2), Linear(n_hidden * 2, n_hidden, bias=False), BatchNorm1D(n_hidden), Tanh(),  
    Linear(n_hidden, vocab_size),
])

Which I think is really cool. I also improved the loss graph, previously it looked like the screenshot on the right, because each iteration’s loss is graphed. Because our batch size is 32, which means in each iteration we look at just 32 inputs/output, a good deal of what the loss will be depends on how lucky we are with the input/output. Over a great number of iterations, it averages out into the correct value, though, so to fix the graph, we instead plot the average of every 1000 iterations, which lets us see a lot more detail that otherwise would be lost in the noise (the two loss graphs I added are not of the same model, but you get the idea). Anyway now I’m going to actually implement the wavenet architecture!