A detailed guide to PyTorch’s nn.Transformer() module. | by Daniel Melchor | Towards Data Science

I have recently been getting more involved in the world of machine learning. When I’ve had a problem understanding a complex issue or coding a neural network, the internet has seemed to have all the answers: from a simple linear regression to complex convolutional networks. At least that is what I thought...
Once I began getting better at this Deep Learning thing, I stumbled upon the all-glorious transformer. The original paper: “Attention is all you need”, proposed an innovative way to construct neural networks. No more convolutions! The paper proposes an encoder-decoder neural network made up of repeated encoder and decoder blocks. The structure is the following:
The left block is the encoder, and the right block is the decoder. If you don’t understand the parts of this model yet, I highly recommend going over Harvard’s “The Annotated Transformer” guide where they code the transformer model in PyTorch from scratch. I will not be covering important concepts like “multi-head attention” or “feed-forward layers” in this tutorial, so you should know them before you continue reading. If you have already taken a look at the code from scratch, you are probably wondering if you are going to have to copy-paste that code all over the place for every project you make. Thankfully, no. Modern python libraries like PyTorch and Tensorflow already include easily accessible transformer models through an import. However, there is more to it than just importing the model and plugging it in. Today I will explain how to use and tune PyTorch nn.Transformer() module. I personally struggled trying to find information about how to implement, train, and infer from it, so I decided to create my own guide for all of you.
To start, we need to import PyTorch and some other libraries we are going to be using:
Now, let’s take a closer look at the transformer module. I recommend starting by reading over PyTorch’s documentation about it. As they explain, there are no mandatory parameters. The module comes with the “Attention is all you need” model hyperparameters. To use it, let’s begin by creating a simple PyTorch model. I will only change some of the default parameters so our model doesn’t take unnecessarily long to train. I made those parameters part of our class:
The transformer blocks don’t care about the order of the input sequence. This, of course, is a problem. Saying “I ate a pizza with pineapple” is not the same as saying “a pineapple ate I with pizza”. Thankfully, we have a solution: positional encoding. This is a way to “give importance” to elements depending on their position. A detailed explanation of how it works can be found here, but a quick explanation is that we create a vector for each element representing its position with regard to every other element in the sequence. Positional encoding follows this very complicated-looking formula which, in practice, we won’t really need to understand:
For the sake of organization and reusability, let’s create a separate class for the positional encoding layer (it looks hard but it is really just the formula, dropout, and a residual connection):
Now that we have the only layer not included in PyTorch, we are ready to finish our model. Before adding the positional encoding, we need an embedding layer so that each element in our sequences is converted into a vector we can manipulate (instead of a fixed integer). We will also need a final linear layer so that we can convert the model’s output into the dimensions of our desired output. The final model should look something like this:
I know... It looks very intimidating, but if you understand what each part does, it is actually a pretty simple model to implement.
You may recall there was a special block in the model structure called “masked multi-head attention”:
So... what is masking? Before I can explain it to you, let’s quickly recapitulate what is going on with our tensors when we feed them into our model. First, we embed and encode (positional encoding) our source tensor. Then, our source tensor is encoded into an unintelligible encoded tensor that we feed into our decoder with our embedded and encoded (positionally) target vector. For our model to learn, we can’t just show it the whole target tensor! This would just give him the answer straight up.
The solution to this is a masking tensor. This tensor is made up of size (sequence length x sequence length) since for every element in the sequence, we show the model one more element. This matrix will be added to our target vector, so the matrix will be made up of zeros in the positions where the transformer can have access to the elements, and minus infinity where it can’t. An illustrated explanation might help you a bit more:
In case you didn’t know, tensors are matrices that can be stored in a GPU, and since they are matrices, all dimensions must have elements of the same size. Of course, this won’t happen when treating with tasks like NLP or different-sized images. Therefore, we use the so-called “special tokens”. These tokens allow our model to know where the start of the sentence is (<SOS>), where the end of the sentence is (<EOS>) and what elements are just there to fill up the remaining space so that our matrices have the sam sequence size (<PAD>). These tokens must also be converted into their corresponding integer id (In our example they will be 2, 3, and 4 respectively). Padding a sequence looks something like this:
To tell our model that these tokens should be irrelevant, we use a binary matrix where there is a True value on the positions where the padding token is and False where it isn’t:
To create the two masking matrices we talked about, we need to extend our transformer model. If you know a bit of NumPy, you will have no problem understanding what these methods do. If you can’t understand it, I recommend opening a Jupyter notebook and going step by step to understand what they do.
The full extended model looks like this (note the change in the forward method as well):
For the sake of this project, I am going to create a set of fake data we can use to train our model. This data will be made up of sequences like:
1, 1, 1, 1, 1, 1, 1, 1 → 1, 1, 1, 1, 1, 1, 1, 1
0, 0, 0, 0, 0, 0, 0, 0 → 0, 0, 0, 0, 0, 0, 0, 0
1, 0, 1, 0, 1, 0, 1, 0 → 1, 0, 1, 0, 1, 0, 1, 0
0, 1, 0, 1, 0, 1, 0, 1 → 0, 1, 0, 1, 0, 1, 0, 1
Feel free to skip to the next section if you aren’t interested in the data creation part.
I won’t bother explaining what these functions do since they are pretty easy to understand with basic NumPy knowledge. I’ll create all the sentences of size 8 so I don’t need padding, and I’ll organize them randomly into batches of size 16:
Now that we have data to work with, we can get to training our model. Let’s begin by creating an instance of our model, loss function, and optimizer. We will use the Stochastic Gradient Descent optimizer, the Cross-Entropy Loss function, and a learning rate of 0.01. I will also use my graphics card for this training since it will take less time, but it is not necessary.
An important concept we need to understand before continuing is that the target tensor we give as an input to the transformer must be shifted by one to the right (compared to the target output tensor). In other words, the tensor we want to give the model for training must have one extra element at the beginning and one less element at the end, and the tensor we compute the loss function with must be shifted in the other direction. This is so that if we give the model an element during inference, it gives us the next one.
Now that we have grasped this concept, let’s get to coding! The training loop is a standard training loop except:
The target tensor is passed to the model during the prediction
A target mask is generated to hide the next words
A padding mask might be generated and passed to the model as well
The validation loop is exactly the same as our training loop except we don’t read or update gradients:
In this example, I am training the model for 10 epochs. To simplify the training I created a fit function that calls the train and validation loop every epoch and prints the loss:
This produces the following out
Training and validating model------------------------- Epoch 1 -------------------------Training loss: 0.5878Validation loss: 0.4172------------------------- Epoch 2 -------------------------Training loss: 0.4384Validation loss: 0.3981------------------------- Epoch 3 -------------------------Training loss: 0.4155Validation loss: 0.3852------------------------- Epoch 4 -------------------------Training loss: 0.4003Validation loss: 0.3700------------------------- Epoch 5 -------------------------Training loss: 0.3842Validation loss: 0.3443------------------------- Epoch 6 -------------------------Training loss: 0.3592Validation loss: 0.3069------------------------- Epoch 7 -------------------------Training loss: 0.3291Validation loss: 0.2652------------------------- Epoch 8 -------------------------Training loss: 0.2956Validation loss: 0.2195------------------------- Epoch 9 -------------------------Training loss: 0.2684Validation loss: 0.1947------------------------- Epoch 10 ------------------------- Training loss: 0.2501Validation loss: 0.1930
After training we obtain the following losses per epoch:
As we can see, our model seems to have learned something. It is time to check if it has, but... how do we check it? We don’t have target tensors for data we have never seen. Here is where shifting our input target and output target tensor has an effect. As we saw before, our model learned to predict the next token when given an element. Therefore, we should be able to give our model the input tensor and the start token, and it should give us back the next element. If when the model predicts a token, we concatenate it with our previous input, we should slowly be able to add words to our output until our model predicts the <EOS> token.
Here is the code for that process:
The output of running this code is:
Example 0Input: [0, 0, 0, 0, 0, 0, 0, 0]Continuation: [0, 0, 0, 0, 0, 0, 0, 0, 0]Example 1Input: [1, 1, 1, 1, 1, 1, 1, 1]Continuation: [1, 1, 1, 1, 1, 1, 1, 1, 1]Example 2Input: [1, 0, 1, 0, 1, 0, 1, 0]Continuation: [1, 0, 1, 0, 1, 0, 1, 0]Example 3Input: [0, 1, 0, 1, 0, 1, 0, 1]Continuation: [1, 0, 1, 0, 1, 0, 1, 0]Example 4Input: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0]Continuation: [0, 1, 0, 1, 0, 1, 0, 1, 0]Example 5Input: [0, 1]Continuation: [0, 1, 0, 1, 0, 1, 0]
So the model has indeed gotten the gist of our sequences, but it still makes some mistakes when trying to predict the continuation. For example, in “Example 4”, the model should predict a 1 as the first token, since the ending of the input is a 0. We can also see how during inference our sentences don’t need to have the same length, and the outputs will also not have the same length (see “Example 5”).
I believe this article can help a lot of beginner/intermediate machine learning developers learn how to work with transformer models in PyTorch, and, since the structure is the same in other languages, this tutorial is probably also useful for other frameworks like Tensorflow (hopefully).
If you have any suggestions or find any bugs feel free to leave a comment and I will fix it ASAP.
Full Colab Notebook: https://drive.google.com/file/d/15gyTrsd-OU6YZVyjwis48ysrXFPEcv9r/view?usp=sharing