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. 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 is complicated in this example, we shall forgo the basic loss criterion used in normal torchbearer models.
/_static/examples/gan.py
Instead use a callback to provide the loss. Since this callback is very simple we can use callback decorators on a function (which takes state) to tell torchbearer when it should be called.
/_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 model on the GPU in the standard way.
/_static/examples/gan.py
We borrow the image saving method from PyTorch_GAN and put it in a call back to save on training step - again using decorators.
/_static/examples/gan.py
After 172400 iterations we see the following.
The source code for the example is given below:
Download Python source code: gan.py </_static/examples/gan.py>