We shall try to implement something more complicated using torchbearer - a Generative Adverserial Network (GAN). This tutorial is a modified version of the GAN from the brilliant collection of GAN implementations PyTorch_GAN by eriklindernoren on github.
We first define all constants for the example.
/_static/examples/gan.py
We then define a number of state keys for convenience using .state_key. This is optional, however, it automatically avoids key conflicts.
/_static/examples/gan.py
We then define the dataset and dataloader - for this example, MNIST.
/_static/examples/gan.py
We use the generator and discriminator from PyTorch_GAN and combine them into a model that performs a single forward pass.
/_static/examples/gan.py
Note that we have to be careful to remove the gradient information from the discriminator after doing the generator forward pass.
Since our loss computation in this example is complicated, we shall forgo the basic loss criterion used in normal torchbearer trials. Instead we use a callback to provide the loss, in this case we use the .add_to_loss callback decorator. This decorates a function that returns a loss and automatically adds this to the main loss in training and validation.
/_static/examples/gan.py
Note that we have summed the separate discriminator and generator losses, since their graphs are separated, this is allowable.
We would like to follow the discriminator and generator losses during training - note that we added these to state during the model definition. We can then create metrics from these by decorating simple state fetcher metrics.
/_static/examples/gan.py
We can then train the torchbearer trial on the GPU in the standard way. Note that when torchbearer is passed a None criterion it automatically sets the base loss to 0.
/_static/examples/gan.py
We borrow the image saving method from PyTorch_GAN and put it in a call back to save ~torchbearer.callbacks.decorators.on_step_training - again using decorators.
/_static/examples/gan.py
Here is a Gif created from the saved images.
The source code for the example is given below:
Download Python source code: gan.py </_static/examples/gan.py>