Fully Convolutional Variational Autoencoder

Question

I would like to make a neural network which uses black and white images as input and outputs a colored version of it. The important thing in that process is that the size of the images must stay the same.

Usually this is done by using a Fully Convolutional Network with GAN or AE architecture. Now I have decided to implement a VAE version, but when i looked it up on the internet I found versions where the latent space was Linear/Dense meaning it breaks the Full Convolution.

Is a VAE not a valid Approach for this type of neural network? Or is there a solution/some code someone can provide which would help in that situation? (Pytorch prefered)

Answer 1

In my experience, when people say "fully" convolutional, even when doing something typical like ImageNet classification, they are still typically referring to a network with at least one final dense layer.

If I understand your question correctly, you're looking to create a VAE with some convolutional layers which has the same sized output as input, but you're confused how to upsample in the decoder such that you go from fewer latent dimensions to an output of the same size as your input. There are a few ways people typically handle this, two of which are deconvolution and sub-pixel (description of the difference, good deconv tutorial). Both allow you to use something like a convolution to take the output from $ d << n $ latent dimensions and upsample it into an $ n $ dimensional output.

Here's an example written in Keras that does this using deconvolutional layers.