VAE/GAN Example

(Code and notebook by TJ Torres, Data Scientist in Data Labs @ StitchFix)

Introduction

This notebook should serve as a brief example showing you how to use the functionality of fauxtograph to simply train a VAE/GAN model on the CLI downloadable hubble image set.

Keep in mind that the performance of the model here represents a low number of training epochs and does not accurately reflect the full potential of the model.

All the code for this can be found on the GitHub fauxtograph repo

Setup

To get setup you'll want to first install the latest version of fauxtograph with

pip install fauxtograph

Or if you have a previous version installed add the upgrade flag as in

pip install --upgrade fauxtograph

With fauxtograph installed you just need to import the model class you'd like (I am importing all three here to show what's availble be we will only use the VAEGAN class) as well as get_paths and image_resize functions. We'll aslop setup the standard inline plotting with matplotlib.

Resizing Images

Since we will be using convolution, and do to the abiguities in upsampling wiht convoluton, the images must be of specific fixed sizes that are determined by the convolution architecture. You can check the next lower allowable size by running:

from fauxtograph import calc_fc_size, calc_im_size

desired_width, desired_height = 100, 100
convolved_shape = calc_fc_size(100, 100)
resize_shape = calc_im_size(*convolved_shape[1:])
print resize_shape

You'll see that we need to resize our images to 96x96.

Next we will get the image paths, then run them through the resizer and regrab those resized paths.

Instantiate and Train the Model

Next we will want to train the model. Let's instantiate it and then use it to load the properly sized images into the variable x_all.

Training the VAE/GAN

Now we will just train the model using the fit method. The options below will save the model and optimizer states as well as some reconstructed/generated images from the model every 2 epochs to the given folders. It'll also print out diagnostics on training losses every epoch and show images generated/reconstructed using the current model state ever pic_freq number of batches. Moreover, the mirroring option randomly mirrors the images along the vertical axis with 50% probability to artificially increase our image set variance. Finally, if you don't have a GPU it'll be best to add the flag_gpu=False option.

NOTE: Be patient with these types of models, they often have somewhat oscillatory behavior at times. Just wait for it to either calm down or make sure you're saving the model state frequently enogh to go back to a former state and retrain.

Saving Your Model

In addition to saving state at given training frequencies, you can also save your model by calling the save method at any point.

Loading a Previous State

It'll then be rather simple to recall the former model state and use it to encode/generate images using the trained model.

Generating Images

Now we can use the previously trained model state to generate images from random normally distributed latent data using the inverse_transform method.

Conclusion

This is just a short introduction but the code is fairly heavily documented with docstrings, so feel free to call the ? operator within iPython and it should help to understand each method and attribute more fully. I hope you enjoy it.