class GAN: from matplotlib import pyplot as plt # draw picture os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices' # enable XLA Devices # define the constants BUFFER_SIZE = 300 ## change the order of image BATCH_SIZE = 1 # batch_size=1 means the trainig is very fast and one picture only trained for 1 time IMG_WIDTH = 256 # define the image size and it must be 256*256*3 because of Unet IMG_HEIGHT = 256 # define the image size LAMBDA = 98 # define the constant number of Lambada OUTPUT_CHANNELS = 3 # channel=1 is enough but 3 can use for color images and it depends on the datasets EPOCHS = 200 # set traning epochs and it is limited by the training sets and Max is 213 because of the size in train folder gan_count=0 "Build the Optimizers of generator and discriminator with Adam method and set the teta_1 as 0.5" optimizer_of_generator_with_Adam_algorithm = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) # optimizer the generator optimizer_of_discriminator_with_Adam_algorithm = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) # optimizer the generator "Build an instance of class Load_Function_Set()" load_function_set = load_function.Load_Function_Set() " load data from train folder and test folder" image_data_in_train_folder = load_function_set.load_and_prepare_data_from_training_folder() # load all images and convert them into suitable data for training in training folder image_data_in_test_folder = load_function_set.load_and_prepare_data_from_test_folder() # load all images and convert them into suitable data for training in test folder "Build Generator with Unet architecture" generator_set = Generator_with_Unet.Generator_Set() generator_with_unet = generator_set.Generator_based_on_paper() # create a generator with U-net "Set the type of calculating loss value as Binary" value_after_Cross_Entropy_calculation = tf.keras.losses.BinaryCrossentropy( from_logits=True) # use binary method to calculate the cross entropy "Build discriminator with VGG architecture" discriminator_set = Discriminator.Discriminator_Set() discriminator = discriminator_set.Discriminator_based_on_VGG() # create discriminator by using downsample "Set the type of calculating loss value as Categorical" value_after_Categorical_Cross_Entropy_calculation = tf.keras.losses.CategoricalCrossentropy # get the loss value by caterorical cross entropy image_processing_set = image_processing.Image_Processing_Set() def __init__(self,name,epoch): self.name=name self.EPOCHS=epoch print("Build "+name+" Successfully") GAN.gan_count+=1 def print_summary_of_G_and_D(self): self.generator_with_unet.summary() # print the summary of generator. Total params: 16,667,907,Trainable params: 16,661,635, Non-trainable params: 6,272 self.discriminator.summary() # output the information of discriminator def plot_model_structure(self): tf.keras.utils.plot_model(self.generator_with_unet, show_shapes=True, dpi=64, to_file='Generator Structure.png') # plot the model of generator tf.keras.utils.plot_model(self.discriminator, show_shapes=True, dpi=64, to_file='Discriminator Structure.png', ) # plot thte discriminator details @tf.function def step_of_training_GAN_with_Gradient(self,corresponding_picture, images_of_real_building, epoch): # define every step trainning with tf.GradientTape() as tape_of_generator, tf.GradientTape() as tape_of_discriminater: # divided the gradient tap into two parts and this is very important output_value_of_generator = self.generator_with_unet(corresponding_picture, training=True) # get the output of generator value_of_disc_for_real_image = self.discriminator([corresponding_picture, images_of_real_building], training=True) # get the disc value of real image value_of_disc_from_generated_image = self.discriminator([corresponding_picture, output_value_of_generator], training=True) # get the disc value of generated image gen_total_loss, gen_gan_loss, gen_l1_loss = calculation_of_loss_value.loss_value_of_generator( value_of_disc_from_generated_image, output_value_of_generator, images_of_real_building) # get the loss value loss_value_of_disc = calculation_of_loss_value.get_the_loss_value_of_discriminator( value_of_disc_for_real_image, value_of_disc_from_generated_image) # get the disc loss value gradients_value_from_discriminator = tape_of_discriminater.gradient(loss_value_of_disc, self. discriminator.trainable_variables) # get the gradients value of discriminator gradients_value_from_generator = tape_of_generator.gradient(gen_total_loss, self. generator_with_unet.trainable_variables) # get the gradients value of generator self.optimizer_of_generator_with_Adam_algorithm.apply_gradients(zip(gradients_value_from_generator, self. generator_with_unet.trainable_variables)) # optimizer generator by using Adam method self.optimizer_of_discriminator_with_Adam_algorithm.apply_gradients( zip(gradients_value_from_discriminator, self. discriminator.trainable_variables)) ##optimizer discriminator Adam method def trainning_GAN(self,dataset_of_trainning, dataset_of_test, epochs): for epoch in range(epochs): # define details in every epoch for an_corresponding_picture, an_target_image in dataset_of_test.take( 1): # for every image in dataset, there will be a predicted image self.image_processing_set.predict_images_from_generator(an_corresponding_picture, an_target_image, self.generator_with_unet) # excuted predict function to generate images print("The image of result is saved in current folder after every 2 epochs") print("Current epoch number: ", epoch," of ", self.name, " is trainning...........") # Train for j, (corresponding_picture, images_of_real_building) in dataset_of_trainning.enumerate(): # apply gradient for iteration and change the generator and discriminator self.step_of_training_GAN_with_Gradient(corresponding_picture, images_of_real_building, epoch) # excute the details of training in every step if epoch % 2 == 0 and epoch > 0: # save the image every 2 epoches self.plt.savefig(str(epoch)+"Epoch" +self.name+ 'result.jpg') # first character is the epoch numbers and if epoch % 200 == 0 and epoch > 0: # show image when the program is finished # plt.show()# show image print("Training is finished") # hint for training finished "Train GAN with the given number of epochs and the data from train folder, test folder " def start_train(self): self.trainning_GAN( self.image_data_in_train_folder, self.image_data_in_test_folder, self.EPOCHS) # Maximize the number of trainingļ¼ but it only has 213 pictures in train sets