For the VAE itself, we instantiate the layers in the __init__ section and define the layer interactions in the ‘forward’ section.
Basically, it’s just a bunch of linear layers with the occasional ReLU activation function.
As VAEs can be a bit finnicky in training (think vanishing and exploding gradients, yay!), I also added two batch normalization layers to the encoder and decoder.
By reducing covariate shift (increasing the independence of two interacting layers) they allow for more stability during training and even have a slight regularization effect (that’s why they get turned off if you switch your model to eval mode).
At the end we need a sigmoid activation function for the binary cross-entropy loss so that all values are between 0 and 1.
The last point which significantly increases training stability is layer initialization.
Inputting a certain set of weights can make a huge difference in training stability, especially for deep neural networks.
I started out with Xavier initialization for all my linear layers which finally allowed me to train the VAE without anything blowing up.
This approach samples initial weights from a random uniform distribution influenced by the number of incoming & outgoing connections of a given layer.
But then recently I stumbled upon an excellent blog post on initialization schemes including Kaiming initialization, so I decided to try that one as well and compare it to the Xavier-initialized VAE.
Apparently, this one is best for ReLU-like activation functions and consists of a weight tensor drawn from a standard normal distribution multiplied by a factor inversely proportional to the number of incoming connections to the layer.
VAE-generated images of blood cells, either healthy (left panel) or infected with malaria (right panel).
Here, the VAE was initialized with Xavier initialization.
Additionally, I added a decaying learning rate (halving after every epoch) to get an improved performance.
After training for 10 epochs, I started to generate images with the trained VAEs.
For this, it’s sufficient to sample random values from the standard normal distribution, input them into the bottleneck layer of the trained VAE and decode them into generated images.
If we have a look at the images generated by the Xavier-initialized VAE, we can clearly see that the VAE did indeed learn something from the images.
Immediately obvious is the color difference between uninfected (yellow) and infected cells (purplish).
Then, if you look closer, uninfected cells seem to be rounder and more evenly shaped than infected cells.
While you can see some granularity in the infected cells it’s not really the same clear-cut clusters as in the input images though.
For images generated from Kaiming-initialized VAEs, we can also observe the clear color difference yet here the granularity seems to be even less pronounced.
Additionally, the images seem quite hazy.
In fact, VAE-generated images have been noted to be kind of noisy.
Sometimes, this isn’t necessarily bad though.
If you remember the ragged edges of the input images, a bit of blurring around the edges at least make the generated images a bit more aesthetically pleasing.
VAE-generated images of blood cells, either healthy (left panel) or infected with malaria (right panel).
Here, the VAE was initialized with Kaiming initialization.
Where to go from here?.Generative adversarial networks (GANs) have been reported to create images with increased resolution in comparison to VAEs, so if that is a factor GANs might be attractive.
Also, especially for images, the usage of linear layers might be inferior to using convolutional layers.
Building a CNN-VAE might yield considerably improved generated images, so go try that if you feel like it!.In principle using an autoencoder for this kind of uninfected / infected cell setup could give you insights about the characteristics of the respective cell states (by investigating the parameters in the constructed latent space) and might help in the automated diagnosis of malaria.
Anyways, I definitely enjoyed working on this little thing and learned a bit more about deep learning & PyTorch.
Looking forward to the next project!Bonus: If you set the parameters right, you can force your VAE to generate cell images that look like gems or colorful pebbles.
Because we all deserve beautiful images!Can you figure out how to get those beautiful cells from a different universe?.