Creating Image Classification Model with Bayesian Perspective
By: Ruben Winastwan
Tags: TensorFlow, Computer Vision, Image Classification, Data Science, Machine Learning

Building Image Classification Model Using Bayesian Neural Networks

In order to build a Bayesian CNN, we need to first import the necessary libraries. In this article, we’re going to use TensorFlow Probability library to create the first and the last layers of our neural networks model.

If you haven’t installed TensorFlow Probability library yet, you can do so by typing the following pip command in your prompt.

pip install tensorflow-probability

Note that you need to upgrade your TensorFlow library to at least version 2.3.0 in order for you to be able to install TensorFlow Probability.

Next, let’s import the necessary libraries to build our neural networks.

Now let’s get started to build our CNN model with a Bayesian perspective. The architecture will be the same as the standard CNN model that we’ve built before. So it has four convolutional layers before we flatten it and use a dense layer at the very end.

You might be wondering, if the CNN architecture will be the same, then what is the difference between the standard CNN model and Bayesian CNN model? Well, the differences are the first convolutional layer and the last dense layer of the model.

The Convolutional2DReparameterization Layer

For the first convolution layer, instead of using a normal CNN Conv2D layer, we are going to use Convolutional2DReparameterization layer. The point of this layer is to take aleatoric uncertainty into account, which is the uncertainty that arises from the quality of the data. So, instead of creating an output of the same deterministic value as a standard CNN does, it will create an output that is drawn from a distribution.

Because this layer creates an output from a distribution, then we need to define several things here as the argument:

Prior for kernel and bias parameters— this is the prior belief on how the distribution looks like before the data is taken into account. In this post, we’re going to define the prior for both kernel and bias parameters with the default normal distribution function from TensorFlow. This normal distribution contains non-trainable parameters, which makes sense for our prior.

kernel_prior_fn = tfpl.default_multivariate_normal_fn

bias_prior_fn = tfpl.default_multivariate_normal_fn,

Posterior for kernel and bias parameters — this is the posterior belief on how the distribution looks like after seeing the evidence from data. As this is a belief after seeing the data, then we need to define the posterior with trainable parameters. For this reason, we can use a default function from TensorFlow for normal distribution.

kernel_posterior_fn = tfpl.default_mean_field_normal_fn(is_singular=False)

bias_posterior_fn = tfpl.default_mean_field_normal_fn(is_singular=False)

Kullback-Leibler divergence — this is the method to measure how one distribution is different from the reference distribution. For our case, this method is used to measure the divergence of our prior and posterior. The lower the value, or when its value is 0, it means that two distributions are coming from the same distribution. If you want to know more in detail about Kullback-Leibler divergence, I will refer you to the following article:

To define this KL divergence, we can use the available function from TensorFlow. We also need to make sure that we scale our KL divergence such that it is only applied once per epoch. A rule of thumb to scale this KL divergence is by dividing it by the total number of your training images.

divergence_fn = lambda q,p,_:tfd.kl_divergence(q,p)/3457 kernel_divergence_fn = divergence_fn

bias_divergence_fn = divergence_fn

Now let’s see how this Convolutional2DReparameterization layer looks like when we initialize it.

divergence_fn = lambda q,p,_:tfd.kl_divergence(q,p)/3457 tfpl.Convolution2DReparameterization(

input_shape=(75,75,3),

filters=8,

kernel_size=16,

activation='relu',

kernel_prior_fn = tfpl.default_multivariate_normal_fn,

kernel_posterior_fn = tfpl.default_mean_field_normal_fn(is_singular=False)

kernel_divergence_fn = divergence_fn,

bias_prior_fn = tfpl.default_multivariate_normal_fn,

bias_posterior_fn = tfpl.default_mean_field_normal_fn(is_singular=False),

bias_divergence_fn = divergence_fn)

As you can see, some of the arguments that we passed on this layer are the same as the standard convolution layer: we need to define the size of the kernel, filter, activation function, and input size. Next, we also add additional arguments for our prior, posterior, and KL divergence as we defined above.

The stack of layers after this layer is the same as the previous standard CNN. What is different is the last dense layer. Instead of using a normal dense layer, we’re going to use the DenseReparameterization layer.

The DenseReparameterization Layer

If the Convolutional2DReparameterization layer in the beginning is to take aleatoric uncertainty into account, this DenseReparameterization layer is applied to take epistemic uncertainty into account. Epistemic uncertainty is the uncertainty that arises from the model itself.

To apply this layer, we need to pass the same arguments as in Convolutional2DReparameterization layer: the kernel prior and posterior, the bias prior and posterior, as well as the divergence function for KL divergence. As an output, we’re going to pass this dense layer to a final one hot categorical layer consists of 5 units, as we have 5 classes in our flowers classification.

tfpl.DenseReparameterization(

units = tfpl.OneHotCategorical.params_size(5),

activation=None,

kernel_prior_fn = tfpl.default_multivariate_normal_fn,

kernel_posterior_fn = tfpl.default_mean_field_normal_fn(is_singular=False)

kernel_divergence_fn = divergence_fn,

bias_prior_fn = tfpl.default_multivariate_normal_fn,

bias_posterior_fn = tfpl.default_mean_field_normal_fn(is_singular=False),

bias_divergence_fn = divergence_fn)

Below is the entire implementation to build our Bayesian CNN model.

Now if you look at the model summary that we have generated from our Bayesian CNN above, we can see that at the first convolutional layer and the last dense layer, the number of parameters is twice as many as the parameters in the standard CNN model.

This is because we have changed the weight of parameters from a single deterministic value into a value that is drawn from a distribution. Since we used a normal distribution, then it has two parameters instead of just one: the mean and the standard deviation.

Compiling and Training the Model

Next, we can compile the model. Same as before, we’re going to use Adam optimizer, and the metrics would be the model’s accuracy. However, we’re going to use a negative log-likelihood as our loss function for our Bayesian CNN model to maximize the likelihood estimate of the mean and the standard deviation of the weights.

Now we are ready to train our model. For this Bayesian CNN, we’re going to train the model for 300 epochs, since the model is more complicated than our standard CNN model.

Predicting Images from Test Data

To predict an image from the test set, we need to define a function to take care all of the necessary steps.

Since we have built a probabilistic model, then we’re not going to get the same probability value when predicting an image. Because of that, we’re going to assign the model to predict our image several times instead of just once. In this post, the model will be assigned to predict a single image 300 times instead of just one time.

Next, based on that 300 predictions, we can create a 95% prediction interval for probabilities. Below is the function to implement that.

Now we can start to predict an image from the test set and see how the model predicts the image. For the predictions, let’s use the same two images that we have used in the standard CNN above: the dandelion one and the rose one.

As you can see above, our model correctly predicts that our image is a dandelion with a very high probability value, which is ~1.0. Not only it predicts our image as dandelion with ~1.0 probability, it also assigns the ~1.0 probability value with an absolute certainty, which is great.

Now let’s take a look at our second image, which is the rose.

With our previous standard CNN, our model predicts correctly that our image is a rose, with 0.999 probability value. That looks great. However, when you look at the result from our Bayesian CNN model above, the aleatoric and epistemic uncertainty actually exist when predicting our rose image:

The aleatoric uncertainty is shown as the model also assigns a high probability value for the dandelion class.

The epistemic uncertainty is shown as the model itself is not certain how big of a probability value that it should assign in each of the flower classes.

As mentioned earlier, we’re assigning the model to predict our rose image 300 times instead of just once. Of those 300 predictions, in one run the model may predict our image as a rose with 0.1 probability and dandelion with 0.8 probability. In another run, it may predict our image as a rose with 1.0 probability and dandelion with 0 probability. The variability of probability value is huge in our CNN model.

With this epistemic uncertainty, our CNN model basically knows what it doesn’t know. It demands us to supply more variations in training data in order for it to classify our rose image with better certainty, like the one with dandelion image. Overall, epistemic uncertainty can be reduced by supplying more training data to the model.