JAX: Guide to Create Convolutional Neural Networks¶

Convolutional Neural Networks (CNN) are a class of neural networks that uses convolution layers applying convolution operations on input data to find out patterns. CNN is commonly used for tasks related to visual imagery (like object detection, image classification, image segmentation, etc.) where they generally output perform other types of neural networks. CNNs are commonly used for natural language processing and time-series as well. CNNs generally has quite less number of parameters to train compared to fully connected dense neural network for tasks with a lot of features. If you want to learn in-depth about convolutional neural networks then please feel free to check our blog which covers theory in detail.

• Convolutional Neural Networks (CNN or ConvNet)

As a part of this tutorial, we'll cover how we can create a simple convolutional neural network using JAX. JAX is a famous framework for designing a neural network that provides functionalities like numpy-like idioms on CPUs/GPUs/TPUs, automatic differentiation, Just-in-time compilation, etc. If readers do not have a background on JAX then we recommend that you go through our tutorial on it that covers basics. It'll help you with this tutorial as well.

• JAX - (Numpy + Automatic Gradients) on Accelerators (GPUs/TPUs)

In this tutorial, we'll be designing a simple convolutional neural network using the high-level stax API of JAX. We have another tutorial on stax API describing how to create simple fully connected neural networks. Please feel free to check it if you are looking for it. It'll help with this tutorial.

• Guide to Create Neural Networks using High-level JAX API

Below we have highlighted important sections of the tutorial.

Important Sections of Tutorial¶

1. Simple Convolutional Neural Network
   • Load Fashion MNIST Dataset
   • Create Neural Network
   • Define Loss Function
   • Train Neural Network (SGD Optimizer)
   • Make Predictions
   • Evaluate Model Performance
   • Train Neural Network (Adam Optimizer)
   • Make Predictions
   • Evaluate Model Performance
2. Guide to Handle Channels First vs Channels Last

Below we have imported JAX and printed the version of it that we'll be using in our tutorial.

Simple Convolutional Neural Network ¶

In this section, we'll explain step by step process to create a convolutional neural network and train it. Below, we have imported the necessary submodules of JAX that we'll be using in our tutorial. We have imported stax and optimizers for creating neural networks and optimizers respectively. The jax.numpy module helps with maintaining arrays.

Load Fashion MNIST Dataset¶

In this section, we have loaded Fashion MNIST dataset available from keras. The dataset has 28x28 size grayscale images of 10 fashion items. The dataset is already divided into the train (60k samples) and test (10k samples) sets. We have converted the dataset loaded as numpy arrays to JAX arrays as required by the model built-in JAX. Later on, we have resized the loaded images to shape (28,28,1) where the last extra dimension represents the channel. The convolutional neural networks require channel dimensions. As our images are grayscale and do not have 3-channels (RGB) like color images, we have introduced an extra dimension as a channel for convolution layers which requires it. After reshaping datasets, we have divided both datasets by float 255 to normalize them as currently images are loaded as integers in the range (0,255). This will help the optimization algorithm converge faster.

Create Neural Network¶

In this section, we have created a convolutional neural network that we'll be using for our fashion MNIST dataset image classification. Our CNN is simple with only 2 convolution layers and one linear/dense layer. The first convolution layer has 32 channels and a kernel size of (3,3). The second layer convolution layer has 16 channels and a kernel size of (3,3). Both have a padding of 'SAME' which indicates that kernel should be applied after padding so that input and output height and width of an image are the same.

The images after first layer will be transformed from shape (n_samples,height,width,1) to (n_samples,height,width,32). The first layer will transform 1 channel image into 32 channels. The height and width in our case are 28. After applying the convolution layer, we have applied Relu (rectified linear unit) activation to the output image. The second layer will transform images from shape (n_samples,height,width,32) to (n_samples,height,width,16). After the second convolution layer, we have again applied Relu activation to the output.

After two convolution layers, we have flattened the output. The new shape now will be (n_samples, height x width x 16) = (n_samples, 28 x 28 x 16) = (n_samples, 12544). After flattening the output, we have added the last dense layer with a number of units the same as our number of classes which is 10. At last, we have applied softmax activation function which maps all 10 output values in the range [0,1] such that the sum of all of them is 1. We can then predict image class as the one with the highest value from that 10.

We have created our neural network using serial() API of stax by giving layers in sequence to it in which they will be applied. The serial() function returns two other functions.

1. init() - This function takes JAX pseudo-random number generator seed and input data shape as input. It then generates model weights and returns them.
2. apply() - This function takes model weights and data as input. It then performs a forward pass of data through the network using weights.

After creating CNN, in the next cell, we have initialized the weights of the neural network using init() method and printed their shapes. Then in the following cell, we have performed a forward pass of a few data samples through the network using apply() method and printed predictions.

Define Loss Function¶

In this section, we have defined the loss function that we'll be using in our case. We'll be using categorical cross entropy function as our loss. The function takes weights, input data, and actual target values as input. The function then makes predictions on input data using weights (apply() method). Then we have one-hot encoded the actual target values as required by our loss function. Then we have taken the log of our predictions. At last, we have multiplied the log of predictions with one-hot encoded actual target values. Then we have taken the sum of the output array.

Later, we'll be calculating the gradient of this loss function with respect to model weights which is the first parameter of the function.

Train Neural Network (SGD Optimizer)¶

In this section, we are actually training our CNN. We have designed a function that takes data and a few other information as input to perform training. The function takes data features, target values, number of epochs, optimizer state (weights), and batch size as input. It then executed the training loop number of epochs time.

For each epoch, it calculates batch indexes for batches of data. It then performs forward pass of batch data through the network to make predictions, calculates loss, and then calculates gradients of loss with respect to weights. It does all that by single call to value_and_grad() function of JAX. This function takes our loss function as input and returns another function. This new function returned by value_and_grad() when called with input values returns two output values. The first value is loss value (cross-entropy) and the second value is gradients of loss with respect to weights. The value_and_grad() function calculates the gradient of the loss with respect to the first parameter of the loss function (cross-entropy) which is weights in our case. The forward pass is performed inside of the loss function.

The weights of our neural network are kept in the optimizer state object. We can retrieve weights of the network by calling get_weights() method of optimizer by giving optimizer state to it. We have explained below the optimizer state when creating the optimizer below.

After calculating gradients, we have updated the weights of the neural network by calling update() method of an optimizer. It takes input gradients and optimizer state (weights). The method returns a new optimizer state which has updated weights.

We are also printing average cross entropy across all batches. At last, we are returning the last updated optimizer state object that has updated weights of the neural network.

Now, we are actually training our neural network by calling a method from the previous cell. We have first initialized the learning rate to 0.0001, a number of epochs to 25, and batch size to 256.

We have then initialized the weights of our CNN using init() method.

Then, we have created SGD optimizer using optimizers module providing learning rate to it. The optimizer returns three functions.

1. init() - It initialized optimizers with weights of neural network and returns initial OptimizerState object which has model weights.
2. update() - It takes gradients and optimizer state as input and returns a new optimizer state which has updated weights (weights updated by subtracting learning rate times gradients).
3. get_weights() - This method takes OptimizerState object as input and returns model weights.

After initializing the optimizer and optimizer state, we have called our function from the previous cell to perform training by providing training data and other parameters. We have stored the final optimizer state returned by a function which we'll use later to make predictions.

We can notice from the cross-entropy loss getting printed that our model seems to be doing a better job. We'll evaluate its performance later by calculating accuracy and classification metrics like precision, recall, and f1-score.

Make Predictions¶

In this section, we are making predictions on train and test data. We have created a small function that loops through data in batches and makes predictions. We have then combined predictions of all batches. The output of our CNN is 10 values per sample hence we have included logic to convert those 10 values to the actual class label. We have predicted the index of maximum value as a class label using argmax() function.

Evaluate Model Performance¶

In this section, we have evaluated the performance of our CNN by calculating the accuracy of train and test predictions. We have used the function available from scikit-learn to calculate the accuracy of train and test predictions.

Then, in the next cell, we have calculated the classification report of test predictions using the function available from scikit-learn. The classification report has information like precision, recall, and f1-score for each target class.

If you want to learn about various metrics available through scikit-learn then please feel free to check our tutorial that covers the majority of them in detail.

• Scikit-Learn - Model Evaluation & Scoring Metrics

Train Neural Network (Adam Optimizer)¶

In this section, we have again trained our CNN but this time using Adam optimizer instead of SGD. The majority of parameters like learning rate, number of epochs, and batch size are the same as previous training. The only change is an optimizer.

Make Predictions¶

In this section, we have made predictions on train and test datasets using new weights calculated using Adam optimizer.

Evaluate Model Performance¶

In this section, we have evaluated the performance of our new CNN trained using Adam optimizer by calculating accuracy and classification report. It seems from the results that in our case SGD optimizer has done a better job compared to Adam.

Guide to Handle Channels First vs Channels Last ¶

Color or RGB images have channels for each color red, green, and blue. Other image types like RGBA, CMYK, HSV can have 3 or more channels that are combined to create the final image. The CNN generally applies kernels on channels of images and transforms them. The images are generally represented as multi-dimensional arrays. To represent a single RGB image, we can use two different multi-dimension formats.

1. Channels First - Here, we maintain 3 2D arrays. The RGB image of shape (28,28) pixel is represented as (3,28,28). There are 3 arrays of shape (28,28).
2. Channels Last - Here, we maintain details of channels for each pixel together. The RGB image of shape (28,28) pixel is represented as (28,28,3). There are 3 values per pixel of (28,28) image.

In our example, we had modified our grayscale images from shape (28,28) to (28,28,1). We had introduced an extra dimension specifying channel for our convolution layers. The default convolution layer Conv() available from JAX requires channel details to be present at last. It can not handle if channel details are present first for images.

Our next two cells explain with examples, how default Conv() layers of JAX properly handle images with channels as the last dimension but fails to handle cases where channels are present first.

The first cell below explains clearly, how it properly transforms an image of shape (28,28) from 1 channel to 16 and then 16 channels to 32 channels. The next cell after it fails to transform channels of input images.

Below, we have printed the definition of Conv() layer and we can notice that there are shape details specified as a tuple of three strings. The first string in the tuple represents input image dimensions, the second string represents kernel dimensions and the third string represents output after applying kernel to input images.

1. NHWC - (Number of samples, Height, Width, Channels) - This is input images dimensions.
2. HWIO - (Height, Width, Input Channels, Output Channels) - This is kernel dimensions of the convolution layer.

We can notice from the default configuration of Conv() layer that it requires channel details present at last.

But, what if we have images with channel details present first. The JAX provides a layer named GeneralConv() that lets us specify the shape of images and kernel as a tuple.

The GeneralConv() layer works exactly like Conv() but the first input to it is a tuple of 3 strings specifying dimensions of input/output images and kernel. We can handle channels first cases using GeneralConv() layer.

Below, we have explained using two simple examples how we can use the GeneralConv() layer when channel details are present first.

This ends our small tutorial explaining how we can create simple convolutional neural networks using JAX. Please feel free to let us know your views in the comments section.

Reference¶

• Convolutional Neural Networks (CNN or ConvNet)
• JAX - (Numpy + Automatic Gradients) on Accelerators (GPUs/TPUs)
• Guide to Create Neural Networks using High-level JAX API
• Guide to Create Simple Neural Networks using JAX