1. First, notice that there is a new optional argument in the Trainer class called model_type, as shown in the following snippet:

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

As you will recall, we used this parameter in Chapter 3, My First GAN in Under 100 Lines, when selecting the MNIST number we wanted to train on. In this instance, we'll use it to select the model type for the generator and discriminator. For this implementation, we will assume that both the generator and discriminator will want to use the same model type. An optional change to consider in the future would be to allow different model types for the generator and discriminator.

2. Next, we need to use the model_type variable with the Discriminator and Generator, as follows:

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

#self.load_MNIST()
self.load_npy(data_path)

The GAN model is unchanged from Chapter 3, My First GAN in Under 100 Lines, at this point, so we simply stick our models in the GAN class to produce a GAN model for adversarial training.

3. The final change worth highlighting in class initialization is data loading; here, we are no longer using clean, built-in datasets. The load_npy function will instead load data from an npy file, as follows:

def load_npy(self,npy_path,amount_of_data=0.25):
        self.X_train = np.load(npy_path)
        self.X_train = 
           self.X_train[:int(amount_of_data*float
                                    (len(self.X_train)))]
        self.X_train = (self.X_train.astype(np.float32) - 
                        127.5)/127.5
        self.X_train = np.expand_dims(self.X_train, axis=3)
        return

Note the following assumptions:

• The data has been cleaned and resized for learning—this means a structure of sample, channels, height, and width.
• The images are all the same size.
• The amount of the data you can ingest is based on your machine; for instance, when using a 1,060 GPU with 16GB of RAM, approximately 50% of the data will load. Remember to modify the amount_of_data parameter based on your own system.