An introduction to unsupervised learning with Fastai
By: Daniel Huynh
Tags: Machine Learning, Vae, Fastai, Variational Autoencoder, Unsupervised Learning

Experiments

Now that we have seen the core spirit of semi supervised learning, where we first exploit unsupervised training without labels, and then train our classifier on top of our feature extractor with all or part of the labels of our samples, we will get our hands dirty with some experiments.

As MNIST has become the “Hello world !” of machine learning, we will also work on this, but this time with a special constraint : we will use only 512 samples from the 60 000 samples of MNIST, so less than 1% of the total data, and use even less labels !

I have chosen to use the excellent framework Fast.ai, to develop a quick and intuitive interface to play with unsupervised learning. For those of you who do not know about Fast.ai I strongly recommend you watching the amazing videos available on Youtube !

To make the training of autoencoders easier, I coded a little library called fastai_autoencoder, which facilitates writing autoencoders, and makes it as clear as the usual fastai code implementations.

Now let’s dive into the code ! The GitHub of the project can be found here : https://github.com/dhuynh95/fastai_autoencoder

First here we import fastai. I will not go into too much detail behind the custom imports from fastai_autoencoder, but to give a quick idea, the bottleneck is the way you chose to filter the information in the encoder. There are several ways to encode the information, which we will cover later. Here we will use a VAE (Variationnal Auto Encoder), which uses a special loss term in addition to the reconstruction loss, a Kullback Divergence term which we will also cover later.

Finally, we import a custom Learner for autoencoders, which is a kind of wrapper to facilitate training, especially to handle callbacks. I created a subclass called VisionAELearner because there are some interesting visualizations that we can get with autoencoders with images.

There is a get_data function which you can have a look at but it’s a bit long, it simply randomly select a few samples from MNIST, like 512, plus some validation example, and store the rest of the samples which will be used as a test set for the classifier.

Here we simply create our data bunch with get_data, which will create a Databunch with empty labels, as we do unsupervised training first, and select only 512 images for training set.

Then we will create our model here. Here you can see that I have tried to factor my code and make it easy to play with, as I have tried a few different options, and tried to implement different papers to see the results. But for now let’s skip this, all you need to know is that at the end, we create a model, by combining three pieces together : the encoder, the bottleneck, and the decoder. Finally we wrap everything in a Fastai learner to use the APIs we are used to.

Notice here that the final size of the bottleneck is 16, which means that we take a 28*28 = 768 pixels image, compress it to 16 variables, and reconstruct a 768 pixels image. This means we have a factor of compression of 48 !

Now we can do the usual learn.lr_find() and learn.recorder.plot(), and have our little plot as usual in Fastai :

Learning rate scheduler

Just before our model starts training, we can have a look at the reconstructions using, a little function I coded called plot_rec(), which is a function of VisionAELearner to plot the reconstructions.

Image reconstruction before training

Here we can see that the reconstruction is just random noise as we have not yet started training. So what are we waiting for ? Let’s go !

Here I proceeded in two phases, first I trained without the KL divergence term, as is suggested here : https://orbit.dtu.dk/files/121765928/1602.02282.pdf

Now that training is over, let us have a look at our reconstructions :

Image reconstruction after training

Ah ! It’s much better, even though we see that there are minor changes, the reconstruction is correct, it even smoothens a little bit the edges, even though we have compressed the size by 48 in the encoder !

To continue having fun, we can see the output of the different steps, using a little function I created, which I call plot_rec_steps, which will give a visualization of the output of the different reconstruction steps done in the decoder through convolutions :

Reconstruction steps in the decoder

Here, the output at step 0 is the output just after we receive our encoded information z of size 16, which has been broadcast to the whole grid using Spatial Broadcast Decoder, which we will cover later. Afterwards we apply several convolutions for the decoder, and we can see here how our image gets reconstructed. First the decoder reconstructs the shape, which can sometimes be similar to other numbers, but then with the last layers the finer details get added and we find the original image.

Finally, we can see bellow the 2D projection using T-SNE of our encoded images z, colored by their classes :

T-SNE projection of encoded images

We can see here, that even though we did not have access to the labels during unsupervised training, the features our autoencoder learn make our data easily separable, as the underlying features generating our images are quite distinct between classes.

Semi supervised training

Now that we have a cool trained autoencoder, we can try training a classifier on top of it and see how it performs. Here we will study how well our model performs in three different settings : when we provide only 128 labels, 256 or all 512 labels from which it had been trained from during the unsupervised training part.

This way, we try to mimic the real settings where you can have only a portion of your images that are labelled, but want to exploit the unlabelled data for training.

One popular technique when working with little data is to use transfer learning. Here I compared using the pretrained encoder from the unsupervised training part, froze it, and added a 2 layers classification head on top of it, and a Resnet18 fine tuned on small datasets of growing size, from 128 labels to 512 labels. We finally compare the accuracy results on the rest of MNIST, which is the 59 000 images that we did not use for unsupervised pretraining. We have the results bellow :

Here we can notice that pretraining the model gives a boost, compared to a Resnet in the very low data regime. Nonetheless, as we increase the amount of labels available, Resnet starts to surpass the pretrained model.

Just for the fun, we can also see how these methods work for Fashion MNIST which is a harder dataset than MNIST, where we also used only 512 images. We first get the following reconstruction :

Reconstruction steps in the decoder for Fashion MNIST

And finally, when retraining a classifier we get the following results :

Comparison of the performances on classification of Fashion MNIST

Conclusion

So in the end, we have quickly seen how unsupervised learning work and how it can be used to pretrained model when there is little data/label. This method is interesting, as it can be used to explore a data set without needing the labels first, and also give a boost when really few labels are present.

In very low label regime it can be interesting, even compared to standard models such as Resnet, though this advantage is lost when more data is present.

For images or text, I do not believe that unsupervised training done the way it was done in this article is competitive compared to traditional transfer learning, but I believe it could be fruitful in other domains, such as signal, tabular, or recommender systems, where it can be useful to use unsupervised learning on your own data set, as there is no big pretrained model available.

I hope you liked this article, and I will come up with other articles soon presenting the theory behind VAE, and how you can implement it in Fastai !

If you have any question, do not hesitate to contact me on Linkedin , you can also find me on Twitter !