Part 2 - The Learning
If you haven't already read the previous blog Part 1- The convolution, in which I discussed the basics of the Convolutions, kernels, etc. In this post, let's discuss the learning part of CNN. How the kernels are updated and the backpropagation of a NN.
For any network when learning, what it basically does is update its weight matrices, in a CNN, the weights are of the kernels/ filters of the convolution layer. This process is done in a process called backpropagation. It involves some mathematics again. Simply put, the weights are updated in a manner that reduces the error of the entire network as close to zero as possible. This error can be simple mean square error or cross entropy loss or some fancy stuff. Which basically tells how bad is the output generated by our network is compared to the ground truth. The kernels are updated such that the NN predicts the correct output every time on unseen data as well.
Let's do some math for better understanding.
In forward propagation, each layer spits out output, which is a result of a matrix multiplication between the kernels and the input matrix (fig 1). The output of the previous layer is fed to the next layer. In the last layer of the network, we get the prediction for the input.
This prediction is used to calculate the error of the network for all the input in the training dataset. Fig. 2 is a simple Mean Absolute Error (MAE) function. Calculates the difference between the actual ground truth and the predicted value and adding over all the data. From this error, we will adjust the values of the kernel, such that in the next iteration, the error is less than that of in the previous iteration.
To update the weights, we adjust the earlier weights with gradients. The gradient which is derived from the error, which we got from the MAE calculations. (more detailed info in this medium article, YouTube). This gradient tells us in which direction the weights need to be updated, with what speed, called the learning rate.
Let's understand more about kernel multiplication and how it is handled by PyTorch.
The first convolution layer in VGG16 : 2-1 Conv2d (3, 64, (3,3)) takes an input image of Channel 3 and returns a matrix of 64 channels. The kernel size is (3x3). But when you look into the weights, it's a tensor of shape [64,3, 3, 3].
Fig 4 shows the math. The convolution (* operator) works on each input channel of the input matrix and the output is an element vice edition of each convolution operation. Which results in output of shape [64,w,h]. K632 stands for 2nd channel of the 63rd filter. [1]
Vgg16:2-3, the second conv2d (64,64, (3,3)) layer, returns output of shape [64, w', h'] and its weights are of shape [64,64,3,3]. Kernel of [3,3] and 64 input channels (from previous layer Conv2d, ReLu), w' and h' are changed size of input because of strides or pooling layer, here they are unchanged.
Now you have a basic understanding of the math behind the convolution. Let's go to fig 3, which is our goal to visualize from the beginning. VGG uses Mean squared error (MSE) instead of MAE and more optimizations.
There are two main problems we face while visualizing CNNs (applies to many ML models).
The tensors, may it be the output or the kernels, all of these are in a higher dimensional space. Even if we could get some tensors to lower dimensions, they might not make any sense. (for us)
For instance, a normal HD image is of shape [3, 1280, 720]. RGB channels. But our weights of kernels are of shape [64,64, 3,3]. 3x3 matrix of 64 channels and such 64 blocks. We somehow need to reduce the dimensions without losing information.
The two methods which I used are linear summation and PCA, I.e. to reduce the shape to [64,1,3,3] with linear summation, to [64,3,3,3] with PCA. From fig 4, we can see that the input channels or the intermediate channels are added to get the output. So.. may be adding all the intermediate kernels/filters may give us something to see. Which will be gray-scale images. If we do a PCA [2] based dimension reduction, we get an RGB image (Down in the videos).
And the learning part? Easy. I captured the values of the kernel at each time stamp, at every epoch. Which will show how they evolve. This is the part that does not make sense. I assumed, during the training process, these kernels evolve and might show some pattern, but with such small [3,3] kernels, it's hard to see any pattern. They look like complete noise. But the output that comes out of looks cool though.
These videos capture the updating kernel weights. Each kernel has different out channels, which are put into rows and columns to make an image. Feature 2 represents all the 64, 3x3 kernels in layer 2 of the VGG16 network. Likewise Feature 7 is layer 7 of VGG16 network (conv2d). I trained the model until the accuracy reached more than 87% on Tiny-ImageNet dataset from Kaggle.
As I said, these look like noise. But one thing we can see is, the kernels in the last layers are changing more than that of the earlier layers. That is because of the vanishing gradient. During the back propagation, gradients become so small, they tend to zero when they reach the starting layer of the network.
Let's look at the output of each channel. These are not captured from scratch, but at the end of training. Passed an image of cat and saved the output after each conv layer. But again the outputs are in higher dimension I did a PCA on each channel. This is how the output looks like.
Original image > output from conv2d layer 0 > 2 > 5 > 7 with PCA
As we go further into the network, the conv layers extract more and more complex features, making the output at each layer more complex and rich with features.