<jupyter_start><jupyter_text>Install gdown Python package.<jupyter_code>!pip install -U --no-cache-dir gdown<jupyter_output><empty_output><jupyter_text>Install Tensorflow-addons.* InstanceNormalization<jupyter_code>!pip install tensorflow-addons<jupyter_output><empty_output><jupyter_text>Install Tensorflow-datasets.* celeb_a dataset<jupyter_code>!pip install tensorflow-datasets<jupyter_output><empty_output><jupyter_text>Import TensorFlow 2.x.<jupyter_code>try:
  %tensorflow_version 2.x
except Exception:
  pass

import tensorflow as tf
import tensorflow.keras.layers as layers
import tensorflow.keras.models as models
from tensorflow.keras import backend as K

import numpy as np
np.random.seed(7)

import matplotlib.pyplot as plot

print(tf.__version__)<jupyter_output><empty_output><jupyter_text>Configuration parameters.<jupyter_code>auto_tune = tf.data.experimental.AUTOTUNE
generator_dimension = 64
generator_residual_blocks = 6
discriminator_dimension = 64
discriminator_residual_blocks = 6
buffer_size = 512
batch_size = 16
epochs = 100
number_of_attributes = 40
image_load_shape = (143, 143, 3)
image_shape = (128, 128, 3)
d_gradient_penalty_weight = 10.0
d_attribute_loss_weight = 1.0
g_attribute_loss_weight = 1.0
g_reconstruction_loss_weight = 10.0
load_previous_weights = False
save_current_weights = False
epsilon = 1e-7<jupyter_output><empty_output><jupyter_text>Compute gradient penalty.<jupyter_code>def gradient_penalty(discriminator, real_image, fake_image):
    sample_shape = [tf.shape(real_image)[0]] + [1] * (real_image.shape.ndims - 1)
    alpha = tf.random.uniform(shape=sample_shape, minval=0., maxval=1.)

    sample_image = real_image + alpha * (fake_image - real_image)
    sample_image.set_shape(real_image.get_shape().as_list())   
    
    with tf.GradientTape(watch_accessed_variables=False) as tape:
        tape.watch(sample_image)

        predictions = discriminator(sample_image, training=False)
        if isinstance(predictions, tuple):
            predictions = predictions[0]

    gradients = tape.gradient(predictions, sample_image)[0]    
    gradients = tf.reshape(gradients, [tf.shape(gradients)[0], -1])
    norm = tf.norm(epsilon + gradients, axis=1)    
    gp_value = tf.reduce_mean((norm - 1.) ** 2)    

    return(gp_value)<jupyter_output><empty_output><jupyter_text>Load CelebA dataset.<jupyter_code>import tensorflow_datasets as tfds
builder = tfds.builder('celeb_a')
print(builder.info)<jupyter_output><empty_output><jupyter_text>Download and prepare dataset.<jupyter_code>#builder.download_and_prepare()<jupyter_output><empty_output><jupyter_text>OR  Download dataset from Google drive.<jupyter_code>from google.colab import drive
drive.mount('/content/drive')
!cp -r '/content/drive/My Drive/datasets/tensorflow_datasets' /root/.<jupyter_output><empty_output><jupyter_text>View dataset contents.<jupyter_code>!ls -al /root/tensorflow_datasets/
!ls -al /root/tensorflow_datasets/celeb_a/<jupyter_output><empty_output><jupyter_text>Load the dataset.<jupyter_code>builder.download_and_prepare()<jupyter_output><empty_output><jupyter_text>Create CelebA dataset splits.* train* validation* test<jupyter_code>celeba_datasets = builder.as_dataset(shuffle_files=True)
print(celeba_datasets)
train_dataset = celeba_datasets['train']<jupyter_output><empty_output><jupyter_text>Preprocess the dataset.<jupyter_code>attributes_to_identifiers = {
    '5_o_Clock_Shadow': 0, 
    'Arched_Eyebrows': 1, 
    'Attractive': 2,       
    'Bags_Under_Eyes': 3,           
    'Bald': 4, 
    'Bangs': 5, 
    'Big_Lips': 6,           
    'Big_Nose': 7, 
    'Black_Hair': 8, 
    'Blond_Hair': 9, 
    'Blurry': 10,           
    'Brown_Hair': 11, 
    'Bushy_Eyebrows': 12, 
    'Chubby': 13,           
    'Double_Chin': 14, 
    'Eyeglasses': 15, 
    'Goatee': 16, 
    'Gray_Hair': 17, 
    'Heavy_Makeup': 18, 
    'High_Cheekbones': 19,          
    'Male': 20, 
    'Mouth_Slightly_Open': 21, 
    'Mustache': 22, 
    'Narrow_Eyes': 23, 
    'No_Beard': 24, 
    'Oval_Face': 25,           
    'Pale_Skin': 26, 
    'Pointy_Nose': 27, 
    'Receding_Hairline': 28,           
    'Rosy_Cheeks': 29, 
    'Sideburns': 30, 
    'Smiling': 31,           
    'Straight_Hair': 32, 
    'Wavy_Hair': 33, 
    'Wearing_Earrings': 34,           
    'Wearing_Hat': 35, 
    'Wearing_Lipstick': 36,           
    'Wearing_Necklace': 37, 
    'Wearing_Necktie': 38, 
    'Young': 39
    }
identifiers_to_attributes = {v: k for k, v in attributes_to_identifiers.items()}<jupyter_output><empty_output><jupyter_text>Create test image.<jupyter_code>def create_test_image(image_shape):
  test_image = np.random.rand(image_shape[0], image_shape[1], image_shape[2])
  return(test_image)<jupyter_output><empty_output><jupyter_text>Create test image batch.<jupyter_code>def create_test_image_batch(image_shape):
  test_image_batch = np.random.rand(batch_size, image_shape[0], image_shape[1], image_shape[2])
  return(test_image_batch)<jupyter_output><empty_output><jupyter_text>Create test attributes.<jupyter_code>def create_test_attributes():
  attributes = np.random.rand(number_of_attributes)
  test_attributes = {}
  for index in range(len(identifiers_to_attributes)):
      attribute = identifiers_to_attributes[index] 
      test_attributes[attribute] = attributes[index]
  return(test_attributes)<jupyter_output><empty_output><jupyter_text>Create test attribute batch.<jupyter_code>def create_test_attribute_batch():
  test_attribute_batch = np.random.rand(batch_size, number_of_attributes)
  return(test_attribute_batch)<jupyter_output><empty_output><jupyter_text>Test created test attributes.<jupyter_code>test_attributes = create_test_attributes()
print(test_attributes)<jupyter_output><empty_output><jupyter_text>Normalize the image to [-1, 1].<jupyter_code>def normalize_image(image):
  image = tf.cast(image, tf.float32)
  image = tf.clip_by_value(image, 0, 255) / 127.5 - 1
  return(image)<jupyter_output><empty_output><jupyter_text>Test image normalization.<jupyter_code>input_image = create_test_image(image_shape)
output_image = normalize_image(input_image)
print('input image shape',input_image.shape)
print('output image shape',output_image.shape)<jupyter_output><empty_output><jupyter_text>Random crop the image.<jupyter_code>def random_crop(image):
  cropped_image = tf.image.random_crop(image, size=image_shape)
  return(cropped_image)<jupyter_output><empty_output><jupyter_text>Test random croping of image.<jupyter_code>input_image = create_test_image(image_load_shape)
output_image = random_crop(input_image)
print('input image shape',input_image.shape)
print('output image shape',output_image.shape)<jupyter_output><empty_output><jupyter_text>Random jitter the image.<jupyter_code>def random_jitter(image):  
  image = tf.image.resize(image, [image_load_shape[0], image_load_shape[1]],
                          method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)  
  image = random_crop(image)
  image = tf.image.random_flip_left_right(image)
  return(image)<jupyter_output><empty_output><jupyter_text>Test random jittering of image.<jupyter_code>input_image = create_test_image(image_load_shape)
output_image = random_jitter(input_image)
print('input image shape',input_image.shape)
print('output image shape',output_image.shape)<jupyter_output><empty_output><jupyter_text>Preprocess train dataset.<jupyter_code>def compute_attributes(attributes_batch):
  attributes_array = []
  for index in range(len(identifiers_to_attributes)):
      attribute = identifiers_to_attributes[index]      
      attributes_array.append(tf.cast(attributes_batch[attribute], dtype=tf.float32))      
    
  return(attributes_array)
def preprocess_train_dataset(sample):
  image = sample['image']
  attributes = sample['attributes']
  
  image = random_jitter(image)
  image = normalize_image(image)

  sample['image'] = image
  sample['attributes'] = compute_attributes(attributes)
  return(sample)<jupyter_output><empty_output><jupyter_text>Test the train dataset preprocessing.<jupyter_code>input_image = create_test_image(image_load_shape)
attributes = create_test_attributes()

sample = {}
sample['image'] = input_image
sample['attributes'] = attributes

output = preprocess_train_dataset(sample)
print('input image shape',input_image.shape)
print('output image shape',output['image'].shape)<jupyter_output><empty_output><jupyter_text>Preprocess the test dataset.<jupyter_code>def preprocess_test_dataset(sample):
  image = sample['image']
  attributes = sample['attributes']

  image = tf.image.resize(image, [image_shape[0], image_shape[1]], method=tf.image.ResizeMethod.NEAREST_NEIGHBOR) 
  image = normalize_image(image)

  sample['image'] = image
  sample['attributes'] = compute_attributes(attributes)
  return(sample)
input_image = create_test_image(image_load_shape)
attributes = create_test_attributes()

sample = {}
sample['image'] = input_image
sample['attributes'] = attributes

output = preprocess_test_dataset(sample)
print('input image shape',input_image.shape)
print('output image shape',output['image'].shape)<jupyter_output><empty_output><jupyter_text>Preprocess dataset splits.<jupyter_code>train_dataset = train_dataset.map(preprocess_train_dataset, num_parallel_calls=auto_tune)
train_dataset = train_dataset.shuffle(buffer_size)
train_dataset = train_dataset.batch(batch_size, drop_remainder=True)
train_dataset = train_dataset.prefetch(auto_tune)<jupyter_output><empty_output><jupyter_text>Create the optimizer.*   Adam optimizer*   Learning rate = 0.0002*   β1 = 0.5*   β2 = 0.999<jupyter_code>generator_optimizer = tf.optimizers.Adam(learning_rate=0.0002, beta_1=0.5, beta_2=0.999)

discriminator_optimizer = tf.optimizers.Adam(learning_rate=0.0002, beta_1=0.5, beta_2=0.999)<jupyter_output><empty_output><jupyter_text>Create the residual block.<jupyter_code>import tensorflow_addons as tfa
def create_residual_block(input_layer, residual_dimension):  

    layer = layers.ZeroPadding2D(padding=1)(input_layer)
    layer = layers.Conv2D(filters=residual_dimension, kernel_size=3, strides=1, padding='valid', use_bias=False)(layer)
    layer = tfa.layers.InstanceNormalization(axis=-1)(layer)
    layer = layers.ReLU()(layer)
    layer = layers.ZeroPadding2D(padding=1)(layer) 
    layer = layers.Conv2D(filters=residual_dimension, kernel_size=3, strides=1, padding='valid', use_bias=False)(layer)
    layer = tfa.layers.InstanceNormalization(axis=-1)(layer)
    return( layers.Add()([input_layer, layer]) )<jupyter_output><empty_output><jupyter_text>Test the residual block.<jupyter_code>residual_dimension = 64 
shape = (128, 128, 64)
input_layer = layers.Input(shape=shape, name='input_attributes') 
residual_block = create_residual_block(input_layer, residual_dimension)

print(residual_block)<jupyter_output><empty_output><jupyter_text>Create the generator model.<jupyter_code>def create_generator_model(image_shape, number_of_attributes=40, generator_dimension=64, generator_residual_blocks=6): 
    input_image = layers.Input(shape=image_shape, name='input_image')
    input_attributes = layers.Input(shape=number_of_attributes, name='input_attributes')        

    attributes = layers.Lambda(lambda x: K.repeat(x, image_shape[0]*image_shape[1]))(input_attributes)
    attributes = layers.Reshape((image_shape[0], image_shape[0], number_of_attributes))(attributes)
    x = layers.Concatenate()([input_image, attributes])
    
    # First Conv2D
    x = layers.Conv2D(filters=generator_dimension, kernel_size=7, strides=1, padding = 'same', use_bias=False)(x)
    x = tfa.layers.InstanceNormalization(axis=-1)(x)
    x = layers.ReLU()(x)
    
    # Downsampling layers
    current_dimension = generator_dimension
    for i in range(2):
      x = layers.ZeroPadding2D(padding=1)(x)
      x = layers.Conv2D(filters = current_dimension*2, kernel_size=4, strides=2, padding = 'valid', use_bias=False)(x)
      x = tfa.layers.InstanceNormalization(axis=-1)(x)
      x = layers.ReLU()(x)
      current_dimension = current_dimension * 2
        
    # Bottleneck layers.
    for i in range(generator_residual_blocks):
      x = create_residual_block(x, current_dimension)
        
    # Upsampling layers
    for i in range(2):
      x = layers.UpSampling2D(size=2)(x)       
      x = layers.Conv2D(filters = current_dimension//2, kernel_size=4, strides=1, padding='same', use_bias=False)(x)
      x = tfa.layers.InstanceNormalization(axis=-1)(x)
      x = layers.ReLU()(x)        
      current_dimension = current_dimension//2
    
    # Last Conv2D
    x = layers.ZeroPadding2D(padding=3)(x)
    generated_image = layers.Conv2D(filters=3, kernel_size=7, strides=1, padding='valid', activation='tanh', use_bias=False, name='generated_image')(x)
    
    return( models.Model(inputs=[input_image, input_attributes], outputs=generated_image, name='generator') )<jupyter_output><empty_output><jupyter_text>Test the generator model.<jupyter_code>generator = create_generator_model(image_shape, number_of_attributes, generator_dimension, generator_residual_blocks)
generator.summary()

input_image = create_test_image_batch(image_shape)
input_attributes = create_test_attribute_batch()

output_image = generator([input_image, input_attributes])
print('input image shape -', input_image.shape)
print('output image shape -', output_image.shape)<jupyter_output><empty_output><jupyter_text>Create the discriminator / classification model.<jupyter_code>def create_discriminator_model(image_shape, number_of_attributes=40, discriminator_dimension=64, discriminator_residual_blocks=6):
    input_image = layers.Input(shape=image_shape, name='input_image')  
    x = layers.ZeroPadding2D(padding=1)(input_image)
    x = layers.Conv2D(filters=discriminator_dimension, kernel_size=4, strides=2, padding='valid', use_bias=False)(x)
    x = layers.LeakyReLU(0.01)(x)
    
    current_dimension = discriminator_dimension
    for i in range(1, discriminator_residual_blocks):
      x = layers.ZeroPadding2D(padding=1)(x)
      x = layers.Conv2D(filters=current_dimension*2, kernel_size=4, strides=2, padding='valid')(x)
      x = layers.LeakyReLU(0.01)(x)
      current_dimension = current_dimension * 2
    
    kernel_size = int(image_shape[0] / np.power(2, discriminator_residual_blocks))    
    
    discriminator_prediction = layers.ZeroPadding2D(padding=1)(x)
    discriminator_prediction = layers.Conv2D(filters=1, kernel_size=4, strides=1, padding='valid', use_bias=False)(discriminator_prediction)        
    discriminator_prediction = layers.Reshape((1, ))(discriminator_prediction)

    discriminator_classification = layers.Conv2D(filters=number_of_attributes, kernel_size=kernel_size, strides=1, padding='valid', use_bias=False)(x)
    discriminator_classification = layers.Reshape((number_of_attributes, ))(discriminator_classification)
    
    return( models.Model(input_image, [discriminator_prediction, discriminator_classification], name='discriminator') )<jupyter_output><empty_output><jupyter_text>Test the discriminator model.<jupyter_code>input_image = create_test_image_batch(image_shape)

discriminator = create_discriminator_model(image_shape, number_of_attributes, discriminator_dimension, discriminator_residual_blocks)
discriminator.summary()
discriminator_prediction, discriminator_classification = discriminator(input_image)

print('discriminator prediction shape', discriminator_prediction.shape)
print('discriminator classification shape', discriminator_classification.shape)<jupyter_output><empty_output><jupyter_text>Create different models and loss functions.* Generator model* Discriminator model* Discriminator loss function* Generator loss function<jupyter_code>generator = create_generator_model(image_shape, number_of_attributes, generator_dimension, generator_residual_blocks)
generator.summary()

discriminator = create_discriminator_model(image_shape, number_of_attributes, discriminator_dimension, discriminator_residual_blocks)
discriminator.summary()<jupyter_output><empty_output><jupyter_text>Load previous model weights.* Generator model weights* Discriminator model weights<jupyter_code>!ls -al '/content/drive/My Drive/models/StarGAN/'
!cp '/content/drive/My Drive/models/StarGAN/'*.h5 .
!ls -al
if(load_previous_weights):
  generator.load_weights('generator.h5')  
  discriminator.load_weights('discriminator.h5')<jupyter_output><empty_output><jupyter_text>Compute the generator loss.<jupyter_code>def classification_loss(Y_true, Y_pred) :
    return tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=Y_true, logits=Y_pred))

def wasserstein_loss(Y_true, Y_pred):
    return K.mean(Y_true*Y_pred)

def reconstruction_loss(Y_true, Y_pred):
    return K.mean(K.abs(Y_true - Y_pred))
def compute_generator_loss(input_image, input_attributes):
  batch_size = input_image.shape[0]
  target_attributes = tf.random.shuffle(input_attributes)

  # Generator
  reconstructed_image = generator([input_image, input_attributes], training=True)
  fake_image = generator([input_image, target_attributes], training=True)  

  # Discriminator
  fake_image_prediction, fake_image_attributes = discriminator(fake_image, training=False)

  valid_predictions = -tf.ones((batch_size, 1))
  fake_image_prediction_loss = wasserstein_loss(valid_predictions, fake_image_prediction)
  fake_image_attributes_loss = classification_loss(target_attributes, fake_image_attributes)    
  image_reconstruction_loss = reconstruction_loss(input_image, reconstructed_image)
   
  generator_loss = (  fake_image_prediction_loss 
                    + fake_image_attributes_loss * g_attribute_loss_weight 
                    + image_reconstruction_loss * g_reconstruction_loss_weight
                    )  

  return(generator_loss)<jupyter_output><empty_output><jupyter_text>Test computation of the generator loss.<jupyter_code>sample_batch = next(iter(train_dataset))
loss = compute_generator_loss(sample_batch['image'], sample_batch['attributes'])
print(loss.numpy())<jupyter_output><empty_output><jupyter_text>Compute the discriminator loss.<jupyter_code>def compute_discriminator_loss(input_image, input_attributes):
  batch_size = input_image.shape[0]
  target_attributes = tf.random.shuffle(input_attributes)

  # Generate
  fake_image = generator([input_image, target_attributes], training=False)    

  # Discriminate
  input_image_prediction, input_image_attributes = discriminator(input_image, training=True)
  fake_image_prediction, fake_image_attributes = discriminator(fake_image, training=True)

  # Discriminator losses
  valid_predictions = -tf.ones((batch_size, 1))
  input_image_gan_loss = wasserstein_loss(valid_predictions, input_image_prediction)

  input_image_attributes_loss = classification_loss(input_attributes, input_image_attributes)

  fake_predictions = tf.ones((batch_size, 1))
  fake_image_gan_loss = wasserstein_loss(fake_predictions, fake_image_prediction)  

  gradient_penalty_value = gradient_penalty(discriminator, input_image, fake_image)        

  discriminator_loss = (  input_image_gan_loss 
                        + fake_image_gan_loss 
                        + gradient_penalty_value * d_gradient_penalty_weight 
                        + input_image_attributes_loss * d_attribute_loss_weight
                        )  
  
  return(discriminator_loss)<jupyter_output><empty_output><jupyter_text>Test computation of the discriminator loss.<jupyter_code>sample_batch = next(iter(train_dataset))
loss = compute_discriminator_loss(sample_batch['image'], sample_batch['attributes'])
print(loss.numpy())<jupyter_output><empty_output><jupyter_text>Train the model.<jupyter_code>model_loss_frequency = 10
model_save_frequency = 50
def train(train_dataset, epochs=100):  
  generator_loss = None
  discriminator_loss = None

  for epoch in range(epochs):
    batch_index = 0
    
    for dataset_batch in train_dataset:  
      batch_index = batch_index + 1    

      images = dataset_batch['image']   
      attributes = dataset_batch['attributes']
      
      if(batch_index%6 == 0):

        with tf.GradientTape(watch_accessed_variables=False) as generator_tape:
          generator_tape.watch(generator.trainable_variables)
          generator_loss = compute_generator_loss(images, attributes)

        generator_gradients = generator_tape.gradient(generator_loss, generator.trainable_variables)        
        generator_optimizer.apply_gradients(zip(generator_gradients, generator.trainable_variables))        

      else:
        with tf.GradientTape(watch_accessed_variables=False) as discriminator_tape:
          discriminator_tape.watch(discriminator.trainable_variables)
          discriminator_loss = compute_discriminator_loss(images, attributes)

        discriminator_gradients = discriminator_tape.gradient(discriminator_loss, discriminator.trainable_variables)
        discriminator_optimizer.apply_gradients(zip(discriminator_gradients, discriminator.trainable_variables))        

      if(batch_index%model_loss_frequency == 0):
        print('generator loss -', generator_loss.numpy(), 'discriminator loss -', discriminator_loss.numpy())

      '''
      if(batch_index%model_save_frequency == 0):
        generator.save_weights('/content/drive/My Drive/models/StarGAN/generator.h5')
        discriminator.save_weights('/content/drive/My Drive/models/StarGAN/discriminator.h5')   
      '''        
train(train_dataset, epochs=10)
if(save_current_weights):
  generator.save_weights('/content/drive/My Drive/models/StarGAN/generator.h5')
  discriminator.save_weights('/content/drive/My Drive/models/StarGAN/discriminator.h5')<jupyter_output><empty_output>
StarcoderdataJupyterStructuredCleanDedup