def convert(): g_gan = generator() f_gan = generator() dis_g = pix2pix.discriminator(norm_type='instancenorm', target=False) dis_f = pix2pix.discriminator(norm_type='instancenorm', target=False) ckpt = tf.train.Checkpoint(ggan=g_gan, fgan=f_gan, gdis=dis_g, fdis=dis_f) if tf.train.latest_checkpoint('checkpoints'): manager = tf.train.CheckpointManager( ckpt, 'checkpoints', 5, ) for index, chpt in enumerate(manager.checkpoints): print(f'{index} : {chpt}') index = int(input('Index')) print('loaded checkpoint:', manager.checkpoints[index]) ckpt.restore(manager.checkpoints[index]) g_gan.save('./g_gan') f_gan.save('./f_gan')
def build_cyclegan_models(n_channels, norm_type): assert norm_type in ['instancenorm', 'batchnorm'] generator_g = pix2pix.unet_generator(n_channels, norm_type=norm_type) generator_f = pix2pix.unet_generator(n_channels, norm_type=norm_type) discriminator_x = pix2pix.discriminator(norm_type=norm_type, target=False) discriminator_y = pix2pix.discriminator(norm_type=norm_type, target=False) return generator_g, generator_f, discriminator_x, discriminator_y
def __init__(self, checkpoint_path: str = None, restore_checkpoint: bool = True): output_channels = 3 logger.info("Creating Generators and Discriminators") self.generator_g = pix2pix.unet_generator( output_channels, norm_type="instancenorm" ) self.generator_f = pix2pix.unet_generator( output_channels, norm_type="instancenorm" ) self.discriminator_x = pix2pix.discriminator( norm_type="instancenorm", target=False ) self.discriminator_y = pix2pix.discriminator( norm_type="instancenorm", target=False ) logger.info("Setting up the optimizers") self.generator_g_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) self.generator_f_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) self.discriminator_x_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) self.discriminator_y_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) self.loss_obj = tf.keras.losses.BinaryCrossentropy(from_logits=True) self.LAMBDA = 10 if checkpoint_path is None: self.checkpoint_path = ".." else: self.checkpoint_path = checkpoint_path self.ckpt = tf.train.Checkpoint( generator_g=self.generator_g, generator_f=self.generator_f, discriminator_x=self.discriminator_x, discriminator_y=self.discriminator_y, generator_g_optimizer=self.generator_g_optimizer, generator_f_optimizer=self.generator_f_optimizer, discriminator_x_optimizer=self.discriminator_x_optimizer, discriminator_y_optimizer=self.discriminator_y_optimizer, ) self.ckpt_manager = tf.train.CheckpointManager( self.ckpt, self.checkpoint_path, max_to_keep=5 ) self.restore_checkpoint = restore_checkpoint if self.restore_checkpoint: # if a checkpoint exists, restore the latest checkpoint. if self.ckpt_manager.latest_checkpoint: self.ckpt.restore(self.ckpt_manager.latest_checkpoint) print("Latest checkpoint restored!!")
def get_models_from_input_shape(input_shape, norm_type, output_init=0.02, residual_output=False): if input_shape == (28, 28, 1): # MNIST-like data return mnist_unet_generator(norm_type=norm_type), \ mnist_discriminator(norm_type=norm_type, target=False) elif input_shape == (256, 256, 3): # TODO: just use our unet_generator fn if residual_output is True or output_init != 0.02: raise NotImplementedError return pix2pix.unet_generator(output_channels=3, norm_type=norm_type), \ pix2pix.discriminator(norm_type=norm_type, target=False) else: return unet_generator(output_channels=3, input_shape=input_shape, norm_type=norm_type, output_init=output_init, residual_output=residual_output), \ pix2pix.discriminator(norm_type=norm_type, target=False)
def __init__(self, summary, lmbda, nsamples, niters, learning_rate, beta_1): self.lmbda = tf.constant(lmbda, tf.float32) self.summary = summary self.ggan = generator() self.fgan = generator() self.gdis = pix2pix.discriminator(norm_type='instancenorm', target=False) self.fdis = pix2pix.discriminator(norm_type='instancenorm', target=False) self.lrscheculer = scheduler.LinearDecay(nsamples * niters // 2, learning_rate, nsamples * niters // 2, 0.0) self.opti_ggan = tf.keras.optimizers.Adam(self.lrscheculer, beta_1) self.opti_gdis = tf.keras.optimizers.Adam(self.lrscheculer, beta_1) self.opti_fgan = tf.keras.optimizers.Adam(self.lrscheculer, beta_1) self.opti_fdis = tf.keras.optimizers.Adam(self.lrscheculer, beta_1)
def __init__(self): build_discriminator = lambda: pix2pix.discriminator(norm_type=NORM_TYPE, target=False) build_optimizer = lambda: tf.keras.optimizers.Adam(2e-4, beta_1=0.5) build_generator = lambda: pix2pix.unet_generator(OUTPUT_CHANNELS, norm_type=NORM_TYPE) # Discriminators (Tries to identify if the input is a real example of its class) self.Dx = build_discriminator() self.Dy = build_discriminator() # Generators (Converts from the opposite class into the new class) self.Gx = build_generator() self.Gy = build_generator() # Optimizers to perform stochastic graident descent self.Dx_optimizer = build_optimizer() self.Dy_optimizer = build_optimizer() self.Gx_optimizer = build_optimizer() self.Gy_optimizer = build_optimizer()
plt.subplot(121) plt.title('Zebra') plt.imshow(sample_zebra[0] * 0.5 + 0.5) plt.subplot(122) plt.title('Zebra with random jitter') plt.imshow(random_jitter(sample_zebra[0]) * 0.5 + 0.5) # Reuse the pix2pix model for the conversion. OUTPUT_CHANNELS = 3 generator_g = pix2pix.unet_generator(OUTPUT_CHANNELS, norm_type='instancenorm') generator_f = pix2pix.unet_generator(OUTPUT_CHANNELS, norm_type='instancenorm') discriminator_x = pix2pix.discriminator(norm_type='instancenorm', target=False) discriminator_y = pix2pix.discriminator(norm_type='instancenorm', target=False) # Plotting the generated images to_zebra = generator_g(sample_horse) to_horse = generator_f(sample_zebra) plt.figure(figsize=(8, 8)) contrast = 8 imgs = [sample_horse, to_zebra, sample_zebra, to_horse] title = ['Horse', 'To Zebra', 'Zebra', 'To Horse'] for i in range(len(imgs)): plt.subplot(2, 2, i + 1) plt.title(title[i])
OUTPUT_CHANNELS = 3 LAMBDA = 10 train_horses = train_horses.map(preprocess_image_train, num_parallel_calls = AUTOTUNE).cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE) train_zebras = train_zebras.map(preprocess_image_train, num_parallel_calls = AUTOTUNE).cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE) test_horses = test_horses.map(preprocess_image_test, num_parallel_calls = AUTOTUNE).cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE) test_zebras = test_zebras.map(preprocess_image_test, num_parallel_calls = AUTOTUNE).cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE) sample_horse = next(iter(train_horses)) sample_zebra = next(iter(train_zebras)) y_generator = pix2pix.unet_generator(OUTPUT_CHANNELS, norm_type = 'instancenorm') x_generator = pix2pix.unet_generator(OUTPUT_CHANNELS, norm_type = 'instancenorm') y_discriminator = pix2pix.discriminator(norm_type = 'instancenorm', target = False) x_discriminator = pix2pix.discriminator(norm_type = 'instancenorm', target = False) generated_y = y_generator(sample_horse) generated_x = x_generator(sample_zebra) prediction_y = y_discriminator(generated_y) prediction_x = x_discriminator(generated_x) y_generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1 = 0.5) x_generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1 = 0.5) y_discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1 = 0.5) x_discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1 = 0.5) EPOCHS = 40
def __init__(self): self.generator = pix2pix.unet_generator(OUTPUT_CHANNELS, norm_type=NORM_TYPE) self.discriminator = pix2pix.discriminator(norm_type=NORM_TYPE, target=False) # Optimizers to perform stochastic graident descent self.generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) self.discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
def preprocess_image_train(image, label): image = random_jitter(image) image = normalize(image) return image def preprocess_image_test(image, label): image = normalize(image) return image train_horses = train_horses.map( preprocess_image_train, num_parallel_calls=AUTOTUNE).cache().shuffle( BUFFER_SIZE).batch(1) train_zebras = train_zebras.map( preprocess_image_train, num_parallel_calls=AUTOTUNE).cache().shuffle( BUFFER_SIZE).batch(1) test_horses = test_horses.map( preprocess_image_test, num_parallel_calls=AUTOTUNE).cache().shuffle( BUFFER_SIZE).batch(1) test_zebras = test_zebras.map( preprocess_image_test, num_parallel_calls=AUTOTUNE).cache().shuffle( BUFFER_SIZE).batch(1) sample_horse = next(iter(train_horses)) sample_zebra = next(iter(train_zebras)) plt.subplot(121) plt.title('Horse') plt.imshow(sample_horse[0] * 0.5 + 0.5) plt.subplot(122) plt.title('Horse with random jitter') plt.imshow(random_jitter(sample_horse[0]) * 0.5 + 0.5) plt.subplot(121) plt.title('Zebra') plt.imshow(sample_zebra[0] * 0.5 + 0.5) plt.subplot(122) plt.title('Zebra with random jitter') plt.imshow(random_jitter(sample_zebra[0]) * 0.5 + 0.5) OUTPUT_CHANNELS = 3 generator_g = pix2pix.unet_generator(OUTPUT_CHANNELS, norm_type='instancenorm') generator_f = pix2pix.unet_generator(OUTPUT_CHANNELS, norm_type='instancenorm') discriminator_x = pix2pix.discriminator(norm_type='instancenorm', target=False) discriminator_y = pix2pix.discriminator(norm_type='instancenorm', target=False) to_zebra = generator_g(sample_horse) to_horse = generator_f(sample_zebra) plt.figure(figsize=(8, 8)) contrast = 8 imgs = [sample_horse, to_zebra, sample_zebra, to_horse] title = ['Horse', 'To Zebra', 'Zebra', 'To Horse'] for i in range(len(imgs)): plt.subplot(2, 2, i+1) plt.title(title[i]) if i % 2 == 0: plt.imshow(imgs[i][0] * 0.5 + 0.5) else: plt.imshow(imgs[i][0] * 0.5 * contrast + 0.5) plt.show() plt.figure(figsize=(8, 8)) plt.subplot(121) plt.title('Is a real zebra?') plt.imshow(discriminator_y(sample_zebra)[0, ..., -1], cmap='RdBu_r') plt.subplot(122) plt.title('Is a real horse?') plt.imshow(discriminator_x(sample_horse)[0, ..., -1], cmap='RdBu_r') plt.show() LAMBDA = 10 loss_obj = tf.keras.losses.BinaryCrossentropy(from_logits=True) def discriminator_loss(real, generated): real_loss = loss_obj(tf.ones_like(real), real) generated_loss = loss_obj(tf.zeros_like(generated), generated) total_disc_loss = real_loss + generated_loss return total_disc_loss * 0.5 def generator_loss(generated): return loss_obj(tf.ones_like(generated), generated) def calc_cycle_loss(real_image, cycled_image): loss1 = tf.reduce_mean(tf.abs(real_image - cycled_image)) return LAMBDA * loss1 def identity_loss(real_image, same_image): loss = tf.reduce_mean(tf.abs(real_image - same_image)) return LAMBDA * 0.5 * loss generator_g_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) generator_f_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) discriminator_x_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) discriminator_y_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) checkpoint_path = "./checkpoints/train" ckpt = tf.train.Checkpoint(generator_g=generator_g, generator_f=generator_f, discriminator_x=discriminator_x, discriminator_y=discriminator_y, generator_g_optimizer=generator_g_optimizer, generator_f_optimizer=generator_f_optimizer, discriminator_x_optimizer=discriminator_x_optimizer, discriminator_y_optimizer=discriminator_y_optimizer) ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=5) # if a checkpoint exists, restore the latest checkpoint. if ckpt_manager.latest_checkpoint: ckpt.restore(ckpt_manager.latest_checkpoint) print ('Latest checkpoint restored!!') EPOCHS = 40 def generate_images(model, test_input): prediction = model(test_input) plt.figure(figsize=(12, 12)) display_list = [test_input[0], prediction[0]] title = ['Input Image', 'Predicted Image'] for i in range(2): plt.subplot(1, 2, i+1) plt.title(title[i]) # getting the pixel values between [0, 1] to plot it. plt.imshow(display_list[i] * 0.5 + 0.5) plt.axis('off') plt.show() @tf.function def train_step(real_x, real_y): # persistent is set to True because the tape is used more than # once to calculate the gradients. with tf.GradientTape(persistent=True) as tape: # Generator G translates X -> Y # Generator F translates Y -> X. fake_y = generator_g(real_x, training=True) cycled_x = generator_f(fake_y, training=True) fake_x = generator_f(real_y, training=True) cycled_y = generator_g(fake_x, training=True) # same_x and same_y are used for identity loss. same_x = generator_f(real_x, training=True) same_y = generator_g(real_y, training=True) disc_real_x = discriminator_x(real_x, training=True) disc_real_y = discriminator_y(real_y, training=True) disc_fake_x = discriminator_x(fake_x, training=True) disc_fake_y = discriminator_y(fake_y, training=True) # calculate the loss gen_g_loss = generator_loss(disc_fake_y) gen_f_loss = generator_loss(disc_fake_x) total_cycle_loss = calc_cycle_loss(real_x, cycled_x) + calc_cycle_loss(real_y, cycled_y) # Total generator loss = adversarial loss + cycle loss total_gen_g_loss = gen_g_loss + total_cycle_loss + identity_loss(real_y, same_y) total_gen_f_loss = gen_f_loss + total_cycle_loss + identity_loss(real_x, same_x) disc_x_loss = discriminator_loss(disc_real_x, disc_fake_x) disc_y_loss = discriminator_loss(disc_real_y, disc_fake_y) # Calculate the gradients for generator and discriminator generator_g_gradients = tape.gradient(total_gen_g_loss, generator_g.trainable_variables) generator_f_gradients = tape.gradient(total_gen_f_loss, generator_f.trainable_variables) discriminator_x_gradients = tape.gradient(disc_x_loss, discriminator_x.trainable_variables) discriminator_y_gradients = tape.gradient(disc_y_loss, discriminator_y.trainable_variables) # Apply the gradients to the optimizer generator_g_optimizer.apply_gradients(zip(generator_g_gradients, generator_g.trainable_variables)) generator_f_optimizer.apply_gradients(zip(generator_f_gradients, generator_f.trainable_variables)) discriminator_x_optimizer.apply_gradients(zip(discriminator_x_gradients, discriminator_x.trainable_variables)) discriminator_y_optimizer.apply_gradients(zip(discriminator_y_gradients, discriminator_y.trainable_variables)) for epoch in range(EPOCHS): start = time.time() n = 0 for image_x, image_y in tf.data.Dataset.zip((train_horses, train_zebras)): train_step(image_x, image_y) if n % 10 == 0: print ('.', end='') n+=1 clear_output(wait=True) # Using a consistent image (sample_horse) so that the progress of the model # is clearly visible. generate_images(generator_g, sample_horse) if (epoch + 1) % 5 == 0: ckpt_save_path = ckpt_manager.save() print ('Saving checkpoint for epoch {} at {}'.format(epoch+1, ckpt_save_path)) print ('Time taken for epoch {} is {} sec\n'.format(epoch + 1, time.time()-start)) # Run the trained model on the test dataset for inp in test_horses.take(5): generate_images(generator_g, inp)
plt.subplot(121) plt.title('Horse') plt.imshow(sample_horse[0] * 0.5 + 0.5) plt.subplot(122) plt.title('Horse with random jitter') plt.imshow(random_jitter(sample_horse[0]) * 0.5 + 0.5) plt.show() OUTPUT_CHANNELS = 3 generator_h_z = pix2pix.unet_generator(OUTPUT_CHANNELS, norm_type='instancenorm') generator_z_h = pix2pix.unet_generator(OUTPUT_CHANNELS, norm_type='instancenorm') discriminator_isHorse = pix2pix.discriminator(norm_type='instancenorm', target=False) discriminator_isZebra = pix2pix.discriminator(norm_type='instancenorm', target=False) to_zebra = generator_h_z(sample_horse) to_horse = generator_z_h(sample_zebra) plt.figure(figsize=(8, 8)) contrast = 8 imgs = [sample_horse, to_zebra, sample_zebra, to_horse] title = ['Horse', 'To Zebra', 'Zebra', 'To Horse'] for i in range(len(imgs)): plt.subplot(2, 2, i+1) plt.title(title[i]) if i % 2 == 0: plt.imshow(imgs[i][0] * 0.5 + 0.5)