The initialization method for this class has two important sections—the variables that we can use to tune the GAN and the initialization of the models themselves. You'll notice that we initialize all of the models and then load the MNIST data in the init method. This ensures that the class is ready to train once it is loaded. A new instance of the class will be needed if you need to change any of the parameters. If those parameters need to be dynamically reconfigurable, you would need to develop a set of helper functions to do that. For the scope of this example, that functionality has been left out.

Here are the steps to build the init class:

1. Initialize the class with the size of the image, latent space, size, number of epochs, batch, and a model_type parameter:

class Trainer:
    def __init__(self, width = 28, height= 28, channels = 1, 
                 latent_size=100, epochs =50000, batch=32, 
                 checkpoint=50,model_type=-1):

2. Make all of the parameters available as internal variables to the class:

self.W = width
self.H = height
self.C = channels
self.EPOCHS = epochs
self.BATCH = batch
self.CHECKPOINT = checkpoint
self.model_type=model_type
self.LATENT_SPACE_SIZE = latent_size

3. Initialize the Generator and Discriminator classes we built in the previous recipes:

self.generator = Generator(height=self.H, width=self.W, 
              channels=self.C, latent_size=self.LATENT_SPACE_SIZE)
self.discriminator = Discriminator(height=self.H, width=self.W,
                                   channels=self.C)
self.gan = GAN(generator=self.generator.Generator, 
               discriminator=self.discriminator.Discriminator)

4. Finally, call the load_MNIST method—this method will pull in the MNIST data into our class automatically:

self.load_MNIST()