Bayesian Neural Networks: 2 Fully Connected in TensorFlow and Pytorch | by Adam Woolf | Towards Data Science

This chapter continues the series on Bayesian deep learning. The previous chapter is available here and the next chapter here. In this chapter we’ll explore alternative solutions to conventional dense neural networks. These alternatives will invoke probability distributions over each weight in the neural network resulting in a single model that effectively contains an infinite ensemble of neural networks trained on the same data. We’ll use this knowledge to solve an important problem of our age: how long to boil an egg.
Become familiar with variational inference with dense Bayesian models
Learn how to convert a normal fully connected (dense) neural network to a Bayesian neural network
Appreciate the advantages and shortcomings of the current implementation
The data is from an experiment in egg boiling. The boil durations are provided along with the egg’s weight in grams and the finding on cutting it open. Findings are categorised into one of three classes: under cooked, soft-boiled and hard-boiled. We want the egg’s outcome from its weight and boiling time. The problem is insanely simple, so much so that the data is near being linearly separable1⁠. But not quite, as the egg’s pre-boil life (fridge temperature or cupboard storage at room temperature) aren’t provided and as you’ll see this swings cooking times. Without the missing data we can’t be certain what we’ll find when opening an egg up. Knowing how certain we are we can influence the outcome here as we can with most problems. In this case if relatively confident an egg’s undercooked we’ll cook it more before cracking it open.
Let’s have a look at the data first to see what we’re dealing with. If you want to feel the difference for yourself you can get the data at github.com/DoctorLoop/BayesianDeepLearning/blob/master/egg_times.csv. You’ll need Pandas and Matplotlib for exploring the data. (pip install — upgrade pandas matplotlib) Download the dataset to the same directory you’re working from. From a Jupyter notebook type pwd on its own in a cell to find out where that directory is if unsure.
And let’s see it now as a histogram.
It seems I wasn’t so good at getting my eggs soft-boiled as I like them so we see a fairly large class imbalance with twice as many underdone instances and three times as many hardboiled instances relative to the soft-boiled lovelies. This class imbalance can spell trouble for conventional neural networks causing them to underperform and an imbalanced class size is a common finding.
Note that we’re not setting density to True (False is the default so doesn’t need to be specified) as we’re interested in comparing actual numbers. While if we were comparing probabilities sampled from one of the three random variables, we’d want to set density=True to normalise the histogram summing the data to 1.0.
Histograms can be troublesome for discretisation of data as we need explicitly specify (or use an algorithm to specify for us) the number of bins to gather values by. Bin size vastly influences how the data appears. As an alternative we can opt to use a kernel density estimate, but it’s often nice when comparing groups as we are here to use a violin plot instead. A violin plot is a hybrid with a box plot showing the smoothed distribution for easy comparison. In python it’s both easier and prettier to do violin plots with the Seaborn library (pip install — upgrade seaborn). But I’ll show it here in matplotlib to be consistent.
Alright great. Now let’s highlight the architectural differences between neural networks we’re actually here to learn about. We’ll implement a classical dense architecture first then transform it into a powerful Bayesianesque equivalent. We implement the dense model with the base library (either TensorFlow or Pytorch) then we use the add on (TensorFlow-Probability or Pyro) to create the Bayesian version. Unfortunately the code for TensorFlow’s implementation of a dense neural network is very different to that of Pytorch so go to the section for the library you want to use.
The source code is available at
github.com/DoctorLoop/BayesianDeepLearning/
We’re using Keras for the implementation as it’s what we need for our later TFP model. We start training by calling the compiled model as shown bellow.
You’ll see TensorBoard is used to keep a record of our training — It’s good to take advantage of TensorBoard as its powerful means of model debugging and one of the key features of TensorFlow. If we want to output training progress beneath the cell during training we can use the Keras output with the verbose parameter of model.fit set to 1. The problem with the Keras output is it’ll print the loss for every epoch. Here we’re training for 800 epochs so we’ll get a long a messy output. It’s therefore nicer to create a custom logger giving control over the frequency of printing the loss. The source code for the custom logger is very simple and available in the notebook at Github. Both the TensorBoard callback (and any custom logger) is passed to the callbacks parameter but this can be ignored if you don’t want to just log every epoch in the notebook. The great thing about Keras is it’ll split the data for validation for us and saving time. Here 10% of the data is used for validation. Some people use 25% of the data for validation but as data is the most important consideration for a neural network and the datasets tend to be fairly large 10% works fine and gives the model a better chance of reaching our training goals. The Github notebook also shows how to use class weights to address the class imbalance we discussed earlier. Finally, ensure validation is infrequent if training time is important as validation is considerably slower than the training itself (especially with a custom logger where approx 10x slower) therefore making training longer.
Final loss is around 0.15 with an accuracy of 0.85. If you train for more than 800 epochs after epoch 1200 you’ll see validation accuracy decreasing even though loss is still declining. That’s because we’ve started to overfit to the training data. When we’re overfitting we’re making our model depend on the specific training instances meaning it doesn’t generalise well to our unseen validation data or the intended real-world application of the model2⁠. Overfitting is the bane of conventional neural networks that often need large datasets and early stopping approaches to mitigate.
But we’re here to learn about solutions not problems! Enter stage our Bayesian neural networks which are resistant to overfitting, aren’t particularly bothered about class imbalance and most importantly perform extraordinarily well with very little data.
Let’s take a look at a Bayesian implementation of the model.
It’s very similar to our conventional model except we’re using a flipout layer from TensorFlow-Probability instead. We specify a kernel divergence function which is the Kullback-Leibler divergence mentioned earlier. That’s about it.
When comparing a conventional dense model to a Bayesian equivalent Pyro does things differently. With Pyro we always create a conventional model first then upgrade it by adding two new functions to make the conversion. The conventional model is needed to provide a way to automatically sample values from the weight distributions. The sampled values are plugged into the corresponding position on the conventional model to produce an estimate of training progress. You’ll remember from the last chapter how training conditions (trains) the weight distributions. In clear speech it means training shifts the normal distributions and alters their scale to represent each weight. However to make a prediction we still plug in solid single-weight values. While we can’t use the distributions wholesale to make predictions we can take many predictions each with different sampled weight values to approximate the distribution. That’s where the dense model we first build fits in nicely.
We’ll start then with a class to enclose our dense model. The full source code is available in an online notebook at https://github.com/DoctorLoop/BayesianDeepLearning.
We’ve setup three layers with only the first two layers utilising an activation function. Then we’ve specified our loss function and optimizer. Be careful to note that for this model we’re utilising the torch.optim.Adam (Pytorch optimizer) rather than the Pyro optimizer. If you try to use Pyro here it’ll throw errors as the parameters are presented differently. Our training loop isn’t anything special if you’ve used Pytorch before. 800 epochs are perfectly excessive, but so what.
[Out]:...Test accuracy: 0.88Final loss:0.173 at epoch: 799 and learning rate: 0.001
Pretty good performance here and quick too. If you’re new to Pytorch and have experience elsewhere you might wonder at the lack of softmax. In Pytorch it’s important to note softmax is built into CrossEntropyLoss.
Let’s get on with upgrading it. Here are the two new functions I mentioned, the model and guide.
We will save intricate discussion about what these function do to another article in the series. Simply put, the model explicitly declares the distributions used for each layer to replace the point values. While the guide declares the variables used to condition (train) those distributions. You’ll notice the functions look similar but on close inspection the model lists the individual distributions for weights and for biases then the guide lists the distributions for the mean and sigma for every weight and bias distribution in the model. That sounds a bit meta — because it is. Save pondering it until later and just have a feel for the training now3.
Training this beautiful monster is fast. We arrive at an accuracy slightly higher than the dense model 0.88–0.90 with a loss of 0.18. I’ve implemented a few extra functions here to make training easier (the trainloader and a predict function) and provide the full source in the GitHub notebook.
We can’t backpropagate through random variables (because by definition they’re random). Therefore we cheat and reparametrize the very distributions and have training update distribution parameters. With this fundamental change we need a different way to calculate training loss.
You may have already known of negative log likelihood that’s used in conventional models. It reflects the probability the data was generated by the model. Don’t worry about what this means now. Instead know that we approximate this loss by taking the average of a large number of samples. But along with our negative log likelihood we combine a new loss leveraging the distribution under calculation. The new loss is Kullback–Leibler divergence (KL divergence) and provides a measure of how different two distributions are from each other. We’ll refer to KL divergence frequently and find it useful in other areas including metrics of uncertainty4. Negative log likelihood is added to the KL divergence to get ELBO loss (expected lower bound on margin likelihood, also known as variational free energy). ELBO loss allows us to approximate distributions during training and benefit hugely in terms of trainability and training time.
Finally, to conclude, let’s glimpse what happens when we make predictions, after all it’s what we’re here for.
We make 5000 samples (this can take a few minutes depending on your computer). With a conventional neural network all these predictions would be exactly the same so it would be a pointless endeavour. But we aren’t using conventional neural networks anymore.
The predict function samples multiple different model versions from the master model we trained with TFP or Pyro. The three test instances we’ve input each result in well defined peaks. The first test instance (red histogram) had a weight of 61.2g and a boiling time of 4.8minutes. Most the time we can see our model predicted it would be a soft-boiled egg, 5% of the time however it predicted an underdone egg and 2% of the time it thought it would be a hard-boiled egg. The model is reasonably confident in this prediction but not as confident as the second test example (green histogram) which almost always predicts a hardboiled result for a 53g egg boiled for 6minutes. How do we achieve consistency from the model if predictions are different each time? We just take the average. If predictions are too variable, i.e. a third in each class we’d not want to make a prediction, except we can act on the information and tell the user that the model is uncertain about the outcome (or ask ourselves if we need work on the model/data some more!) In both cases we gain really powerful information about something a conventional model can’t offer. We can therefore think of a conventional model as always arrogant while our new model is appropriate. It’s cautious, but confident when confidence is correct. We don’t need be a computer-science psychologist to appreciate the preferable model-personality.
In this chapter we’ve explored the changes in a basic Bayesian model and seen some major advantages accompanying it. We’ve seen how to do this in TensorFlow-Probability and in Pyro. While the model is fully functional, at this stage it isn’t perfect and neither is it truly Bayesian. Whether a ‘truer’ model matters depends on your circumstances. In subsequent articles we’ll include discussion about imperfections like the use of softmax and how we can address them. In the next article our main focus will be on image prediction with Bayesian convolutional deep learning. If my writing looked a little more formal I’d perhaps be more confident of seeing you there!
1 Therefore it’s somewhat of an abuse to solve the problem with a neural network when simpler models would do fine.
2 Unsure what the real-world application is here, a breakfast cafe?
3 If your thinking: Wow the TensorFlow-Probability code looked way simpler it’s because it does the operations for you for now. That said, we can make the Pyro code simpler as well by using a Pyro autoguide. As it suggests this plays the role of the guide for you. But we’re here to learn the wonder of Bayesian deep learning so we need to get exposure at some point!
4 Interestingly we can make a simple classifier with reasonable performance by using KL divergence directly, i.e. an image pattern classifier can be built with the local binary pattern algorithm and KL. In this partnership KL compares the distributions of corners and edges between the input image and an image with a known pattern.