Implementing Yann LeCun’s LeNet-5 in PyTorch | by Eryk Lewinson | Towards Data Science

Recently, I watched the Data Science Pioneers movie by Dataiku, in which several data scientists talked about their jobs and how they apply data science in their daily jobs. In one of the talks, they mention how Yann LeCun’s Convolutional Neural Network architecture (also known as LeNet-5) was used by the American Post office to automatically identify handwritten zip code numbers. Another real-world application of the architecture was recognizing the numbers written on cheques by banking systems. It makes sense to point out that the LeNet-5 paper was published in 1998. That is one of the reasons why it is a good starting point to understand how CNNs work, before moving to more complex and modern architectures.
In this article, I briefly describe the architecture and show how to implement LeNet-5 in PyTorch. The article assumes a general understanding of the basics of Convolutional Neural Networks, including concepts such as convolutional layers, pooling layers, fully-connected layers, etc.
LeNet-5 is a 7 layer Convolutional Neural Network, trained on grayscale images of size 32 x 32 pixels.
To decode the image above, I present the naming convention used by the authors:
Cx — convolution layer,
Sx — subsampling (pooling) layer,
Fx — fully-connected layer,
x — index of the layer.
As LeNet-5 is relatively simple given modern standards, we can go over each of the layers separately to gain a good understanding of the architecture. However, before proceeding it also makes sense to remind the formula for calculating the output size of the convolutional layer.
The formula is presented below:
, where W is the input height/width (normally the images are squares, so there is no need to differentiate the two), F is the filter/kernel size, P is the padding, and S is the stride.
Having seen the architecture schema and the formula above, we can go over each layer of LeNet-5.
Layer 1 (C1): The first convolutional layer with 6 kernels of size 5×5 and the stride of 1. Given the input size (32×32×1), the output of this layer is of size 28×28×6.Layer 2 (S2): A subsampling/pooling layer with 6 kernels of size 2×2 and the stride of 2. The subsampling layer in the original architecture was a bit more complex than the traditionally used max/average pooling layers. I will quote [1]: “ The four inputs to a unit in S2 are added, then multiplied by a trainable coefficient, and added to a trainable bias. The result is passed through a sigmoidal function.”. As a result of non-overlapping receptive fields, the input to this layer is halved in size (14×14×6).Layer 3 (C3): The second convolutional layer with the same configuration as the first one, however, this time with 16 filters. The output of this layer is 10×10×16.Layer 4 (S4): The second pooling layer. The logic is identical to the previous one, but this time the layer has 16 filters. The output of this layer is of size 5×5×16.Layer 5 (C5): The last convolutional layer with 120 5×5 kernels. Given that the input to this layer is of size 5×5×16 and the kernels are of size 5×5, the output is 1×1×120. As a result, layers S4 and C5 are fully-connected. That is also why in some implementations of LeNet-5 actually use a fully-connected layer instead of the convolutional one as the 5th layer. The reason for keeping this layer as a convolutional one is the fact that if the input to the network is larger than the one used in [1] (the initial input, so 32×32 in this case), this layer will not be a fully-connected one, as the output of each kernel will not be 1×1.Layer 6 (F6): The first fully-connected layer, which takes the input of 120 units and returns 84 units. In the original paper, the authors used a custom activation function — a variant of the tanh activation function. For a thorough explanation, please refer to Appendix A in [1].Layer 7 (F7): The last dense layer, which outputs 10 units. In [1], the authors used Euclidean Radial Basis Function neurons as activation functions for this layer.
Layer 1 (C1): The first convolutional layer with 6 kernels of size 5×5 and the stride of 1. Given the input size (32×32×1), the output of this layer is of size 28×28×6.
Layer 2 (S2): A subsampling/pooling layer with 6 kernels of size 2×2 and the stride of 2. The subsampling layer in the original architecture was a bit more complex than the traditionally used max/average pooling layers. I will quote [1]: “ The four inputs to a unit in S2 are added, then multiplied by a trainable coefficient, and added to a trainable bias. The result is passed through a sigmoidal function.”. As a result of non-overlapping receptive fields, the input to this layer is halved in size (14×14×6).
Layer 3 (C3): The second convolutional layer with the same configuration as the first one, however, this time with 16 filters. The output of this layer is 10×10×16.
Layer 4 (S4): The second pooling layer. The logic is identical to the previous one, but this time the layer has 16 filters. The output of this layer is of size 5×5×16.
Layer 5 (C5): The last convolutional layer with 120 5×5 kernels. Given that the input to this layer is of size 5×5×16 and the kernels are of size 5×5, the output is 1×1×120. As a result, layers S4 and C5 are fully-connected. That is also why in some implementations of LeNet-5 actually use a fully-connected layer instead of the convolutional one as the 5th layer. The reason for keeping this layer as a convolutional one is the fact that if the input to the network is larger than the one used in [1] (the initial input, so 32×32 in this case), this layer will not be a fully-connected one, as the output of each kernel will not be 1×1.
Layer 6 (F6): The first fully-connected layer, which takes the input of 120 units and returns 84 units. In the original paper, the authors used a custom activation function — a variant of the tanh activation function. For a thorough explanation, please refer to Appendix A in [1].
Layer 7 (F7): The last dense layer, which outputs 10 units. In [1], the authors used Euclidean Radial Basis Function neurons as activation functions for this layer.
For more details on the reasoning behind the architecture and some choices (such as the non-standard activation functions), please refer to [1].
I will now show how to implement LeNet-5 (with some minor simplifications) in PyTorch. To keep the spirit of the original application of LeNet-5, we will train the network on the MNIST dataset. We start by importing all the required libraries.
Additionally, we check if the GPU is available and set the DEVICE variable accordingly. In the next step, we set up some parameters (such as the random seed, learning rate, batch size, number of epochs, etc.), which we will use later on while setting up the neural network.
As the next step, we define some helper functions used for training the neural network in PyTorch. As the general idea is very similar in most cases, you can slightly modify the functions for your needs and use them for training all kinds of networks. We start with the function responsible for the training part:
I will quickly describe what is happening in the train function. For each batch of observations we carry out the following steps:
Perform the forward pass — getting the predictions for the batch using the current weightsCalculate the value of the loss functionPerform the backward pass, in which the weights are adjusted based on the loss. This is the “learning” step.
Perform the forward pass — getting the predictions for the batch using the current weights
Calculate the value of the loss function
Perform the backward pass, in which the weights are adjusted based on the loss. This is the “learning” step.
Please note that for the training phase, the model is in the training mode (model.train()) and we also need to zero out the gradients for each batch. Additionally, we calculate the running loss within the training step.
Then, we define the function responsible for validation.
The validation function is very similar to the training one, with the difference being the lack of the actual learning step (the backward pass). Please note that we need to specify that we are using the model for evaluation only — model.eval(). We do not need to worry about the gradients, as in the next function you will see that we disable them for the validation step. Lastly, we combine them all together within the training loop:
In the training loop, for each epoch, we run both the train and validate functions, with the latter one running with torch.no_grad() in order not to update weights and save some computation time. Additionally to the loss function used for training, we calculate the accuracy of the model for both the training and validation steps using the custom get_accuracy function. The dataset we will be using is balanced, so there is no problem with using accuracy as the metric of choice. For brevity, I do not include all the helper functions here and you can find the definition of get_accuracy and plot_losses on GitHub.
Having defined the helper functions, it is time to prepare the data. Thanks to the popularity of the MNIST dataset (if you are not familiar with it, you can read some background here), it is readily available as one of the datasets within torchvision.
In the snippet above, we first defined a set of transforms to be applied to the source images. We first resize the images to 32×32 (the input size of LeNet-5) and then convert them to tensors. transforms.ToTensor() automatically scales the images to the [0, 1] range. Instead of resizing the images, we could have also applied some sort of padding to the images. In the simplest case, we would simply add 2 zeros on each side of the original image.
The second step is to define the datasets. For the training object, we specified download=True in order to download the dataset. While defining the datasets, we also indicate the previously defined transformations and whether the particular object will be used for training or not.
Lastly, we instantiated the DataLoaders by providing the dataset, the batch size, and the desire to shuffle the dataset in each epoch. For validation, this does not make a difference so we set it to False.
Below you can see a preview of 50 images coming from the training set.
Finally, it is the time to define the LeNet-5 architecture.
From the class definition above, you can see a few simplifications in comparison to the original network:
using the regular tanh activation function instead of the custom one described in [1],
using average pooling layers instead of the more complex equivalents used in the original architecture,
replacing the Euclidean Radial Basis Function activations in the output layer with the softmax function.
After defining the class, we need to instantiate the model (and send it to the correct device), the optimizer (ADAM in this case), and the loss function (Cross entropy). The last two are different than the method originally used in [1].
Having everything in place, we can train the network by running:
model, optimizer, _ = training_loop(model, criterion, optimizer, train_loader, valid_loader, N_EPOCHS, DEVICE)
What results in the following outputs:
The training loss plateaus, while the validation loss sometimes exhibits small bumps (increased values). Overall, I believe the performance can be described as quite satisfactory. The best results (on the validation set) were achieved in the 11th epoch.
To evaluate the predictions of our model, we can run the following code which displays a set of numbers coming from the validation set, together with the predicted label and the probability that the network assigns to that label (in other words, how confident the network is in the prediction).
From the image above we can see that the network is almost always sure about the label, with the only doubt visible in the digit 3 (in the second row, 3rd from right), when it is only 54% sure it is a 3. That is most likely due to the fact that the number indeed resembles an 8.
As the very last step, we run the following command to remove the downloaded dataset:
!rm -r mnist_data
In this article, I described the architecture of LeNet-5 and showed how to implement it and train it using the famous MNIST dataset. As this is one of the first CNN architectures, it is relatively straightforward and easy to understand, which makes it a good start for learning about Convolutional Neural Networks.
To further improve the performance of the network, it might be worthwhile to experiment with some data augmentation. To do so, we can apply transformations such as rotation or shearing (using torchvision.transforms) to the image, in order to create a more varied dataset. We should also pay attention that not all transformations might be applicable to the case of digit recognition. An example of such an incorrect transformation would be flipping the image to create a mirror reflection.
You can find the code used for this article on my GitHub. As always, any constructive feedback is welcome. You can reach out to me on Twitter or in the comments.
[1] Y. LeCun, L. Bottou, Y. Bengio, and P. Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, November 1998. — available here