{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"<center><img src='https://raw.githubusercontent.com/dimitreOliveira/MachineLearning/master/Kaggle/I%E2%80%99m%20Something%20of%20a%20Painter%20Myself/banner.png' height=350></center>\n<p>\n<h1><center> I’m Something of a Painter Myself </center></h1>\n<h2><center> Improving CycleGAN - Monet paintings </center></h2>\n\n\n### This notebook was created to try several experiments to the original CycleGan (or at least a vey close one) implementation, and see what works or not for this specific task.\n\n#### Experiments (Starting from the original paper architecture):\n- Transformer with residual blocks [++] \n- Residual connections between Generator and Discriminator [++]\n- Not using `InstanceNorm` at the first layer of both generator and discriminator [++]\n- Better `InstanceNorm` layer initialization [++]\n- Training a lot longer [++]\n- Better `Conv` layer initialization [+]\n- Residual connection with `Concatenate` instead of `Add` [+]\n- Data augmentations (flips, rotations, and crops) [+]\n- Discriminator with label smoothing [+]\n- Using [external data](https://www.kaggle.com/dimitreoliveira/tfrecords-monet-paintings-256x256) (1193 files) [+-]\n- Train on crops [+-] \n- Decoder with resize-convolution [+-]\n- 9 transformer blocks [+-] \n- Patch discriminator [+-] \n- Lager batch size [+-] \n\n\n#### Some references\n- [My previous work with a CycleGAN introduction and baseline](https://www.kaggle.com/dimitreoliveira/introduction-to-cyclegan-monet-paintings)\n- [CycleGAN homepage](https://junyanz.github.io/CycleGAN/)\n- [Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks - Original paper](https://arxiv.org/pdf/1703.10593.pdf)\n- [CycleGAN: Learning to Translate Images (Without Paired Training Data)](https://towardsdatascience.com/cyclegan-learning-to-translate-images-without-paired-training-data-5b4e93862c8d)\n- [One nice implementation of TF CycleGAN](https://github.com/architrathore/CycleGAN)\n- [How to Develop a CycleGAN for Image-to-Image Translation with Keras](https://machinelearningmastery.com/cyclegan-tutorial-with-keras/)\n- [How to Implement GAN Hacks in Keras to Train Stable Models](https://machinelearningmastery.com/how-to-code-generative-adversarial-network-hacks/)\n- [How to Train a GAN? Tips and tricks to make GANs work](https://github.com/soumith/ganhacks)\n- [Deconvolution and Checkerboard Artifacts](https://distill.pub/2016/deconv-checkerboard/)\n- [Keras implementation of CycleGAN](https://keras.io/examples/generative/cyclegan/)\n- [CycleGAN with Better Cycles](https://ssnl.github.io/better_cycles/report.pdf)","metadata":{}},{"cell_type":"markdown","source":"## Dependencies","metadata":{}},{"cell_type":"code","source":"import os, random, json, PIL, shutil, re, imageio, glob\nimport numpy as np\nimport pandas as pd\nimport seaborn as sns\nfrom PIL import ImageDraw\nimport matplotlib.pyplot as plt\nfrom kaggle_datasets import KaggleDatasets\nimport tensorflow as tf\nimport tensorflow.keras.layers as L\nimport tensorflow.keras.backend as K\nimport tensorflow_addons as tfa\nfrom tensorflow.keras import Model, losses, optimizers\nfrom tensorflow.keras.callbacks import Callback\n\n\ndef seed_everything(seed=0):\n    random.seed(seed)\n    np.random.seed(seed)\n    tf.random.set_seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    os.environ['TF_DETERMINISTIC_OPS'] = '1'\n    \nSEED = 0\nseed_everything(SEED)","metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## TPU configuration","metadata":{}},{"cell_type":"code","source":"try:\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n    print(f'Running on TPU {tpu.master()}')\nexcept ValueError:\n    tpu = None\n\nif tpu:\n    tf.config.experimental_connect_to_cluster(tpu)\n    tf.tpu.experimental.initialize_tpu_system(tpu)\n    strategy = tf.distribute.experimental.TPUStrategy(tpu)\nelse:\n    strategy = tf.distribute.get_strategy()\n\n\nREPLICAS = strategy.num_replicas_in_sync\nAUTO = tf.data.experimental.AUTOTUNE\nprint(f'REPLICAS: {REPLICAS}')","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model parameters","metadata":{}},{"cell_type":"code","source":"HEIGHT = 256\nWIDTH = 256\nHEIGHT_RESIZE = 128\nWIDTH_RESIZE = 128\nCHANNELS = 3\nBATCH_SIZE = 16\nEPOCHS = 120\nTRANSFORMER_BLOCKS = 6\nGENERATOR_LR = 2e-4\nDISCRIMINATOR_LR = 2e-4","metadata":{"_kg_hide-input":false,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Load data","metadata":{}},{"cell_type":"code","source":"GCS_PATH = KaggleDatasets().get_gcs_path('gan-getting-started')\n# EXT_PATH = KaggleDatasets().get_gcs_path('tfrecords-monet-paintings-256x256')\n# EXT_PATH = KaggleDatasets().get_gcs_path('monet-tfrecords-256x256')\n\nMONET_FILENAMES = tf.io.gfile.glob(str(GCS_PATH + '/monet_tfrec/monet*.tfrec'))\nPHOTO_FILENAMES = tf.io.gfile.glob(str(GCS_PATH + '/photo_tfrec/photo*.tfrec'))\n# MONET_FILENAMES = tf.io.gfile.glob(str(EXT_PATH + '/monet*.tfrec'))\n\ndef count_data_items(filenames):\n    n = [int(re.compile(r\"-([0-9]*)\\.\").search(filename).group(1)) for filename in filenames]\n    return np.sum(n)\n\nn_monet_samples = count_data_items(MONET_FILENAMES)\nn_photo_samples = count_data_items(PHOTO_FILENAMES)\n\nprint(f'Monet TFRecord files: {len(MONET_FILENAMES)}')\nprint(f'Monet image files: {n_monet_samples}')\nprint(f'Photo TFRecord files: {len(PHOTO_FILENAMES)}')\nprint(f'Photo image files: {n_photo_samples}')","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Augmentations\n\nData augmentation for GANs should be done very carefully, especially for tasks similar to style transfer, if we apply transformations that can change too much the style of the data (e.g. brightness, contrast, saturation) it can cause the generator to do not efficiently learn the base style, so in this case, we are using only spatial transformations like, flips, rotates and crops.","metadata":{}},{"cell_type":"code","source":"def data_augment(image):\n    p_spatial = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    p_rotate = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n#     p_crop = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n\n    \n#     # Random jitter\n#     image = tf.image.resize(image, [286, 286], method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)\n    \n    # 90º rotations\n    if p_rotate > .8:\n        image = tf.image.rot90(image, k=3) # rotate 270º\n    elif p_rotate > .6:\n        image = tf.image.rot90(image, k=2) # rotate 180º\n    elif p_rotate > .4:\n        image = tf.image.rot90(image, k=1) # rotate 90º\n        \n    # Flips\n    image = tf.image.random_flip_left_right(image)\n    image = tf.image.random_flip_up_down(image)\n    if p_spatial > .75:\n        image = tf.image.transpose(image)\n        \n#     # Crops\n#     if p_crop > .6: # random crop\n#         crop_size = tf.random.uniform([], int(HEIGHT*.7), HEIGHT, dtype=tf.int32)\n#         image = tf.image.random_crop(image, size=[crop_size, crop_size, CHANNELS])\n#     elif p_crop > .2: # central crop\n#         if p_crop > .5:\n#             image = tf.image.central_crop(image, central_fraction=.7)\n#         elif p_crop > .35:\n#             image = tf.image.central_crop(image, central_fraction=.8)\n#         else:\n#             image = tf.image.central_crop(image, central_fraction=.9)\n            \n    # Train on crops\n    image = tf.image.random_crop(image, size=[HEIGHT_RESIZE, WIDTH_RESIZE, CHANNELS])\n        \n    \n    return image","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Auxiliar functions","metadata":{}},{"cell_type":"code","source":"def normalize_img(img):\n    img = tf.cast(img, dtype=tf.float32)\n    # Map values in the range [-1, 1]\n    return (img / 127.5) - 1.0\n\ndef decode_image(image):\n    image = tf.image.decode_jpeg(image, channels=CHANNELS)\n    image = tf.reshape(image, [HEIGHT, WIDTH, CHANNELS])\n    return image\n\ndef read_tfrecord(example):\n    tfrecord_format = {\n        'image':      tf.io.FixedLenFeature([], tf.string)\n    }\n    example = tf.io.parse_single_example(example, tfrecord_format)\n    image = decode_image(example['image'])\n    return image\n\ndef load_dataset(filenames):\n    dataset = tf.data.TFRecordDataset(filenames)\n    dataset = dataset.map(read_tfrecord, num_parallel_calls=AUTO)\n    return dataset\n\ndef get_dataset(filenames, augment=None, repeat=True, shuffle=True, batch_size=1):\n    dataset = load_dataset(filenames)\n\n    if augment:\n        dataset = dataset.map(augment, num_parallel_calls=AUTO)\n    dataset = dataset.map(normalize_img, num_parallel_calls=AUTO)\n    if repeat:\n        dataset = dataset.repeat()\n    if shuffle:\n        dataset = dataset.shuffle(512)\n        \n    dataset = dataset.batch(batch_size)\n    dataset = dataset.cache()\n    dataset = dataset.prefetch(AUTO)\n    \n    return dataset\n\ndef display_samples(ds, n_samples):\n    ds_iter = iter(ds)\n    for n_sample in range(n_samples):\n        example_sample = next(ds_iter)\n        plt.subplot(121)\n        plt.imshow(example_sample[0] * 0.5 + 0.5)\n        plt.axis('off')\n        plt.show()\n        \ndef display_generated_samples(ds, model, n_samples):\n    ds_iter = iter(ds)\n    for n_sample in range(n_samples):\n        example_sample = next(ds_iter)\n        generated_sample = model.predict(example_sample)\n        \n        f = plt.figure(figsize=(12, 12))\n        \n        plt.subplot(121)\n        plt.title('Input image')\n        plt.imshow(example_sample[0] * 0.5 + 0.5)\n        plt.axis('off')\n        \n        plt.subplot(122)\n        plt.title('Generated image')\n        plt.imshow(generated_sample[0] * 0.5 + 0.5)\n        plt.axis('off')\n        plt.show()\n        \ndef evaluate_cycle(ds, generator_a, generator_b, n_samples=1):\n    fig, axes = plt.subplots(n_samples, 3, figsize=(22, (n_samples*6)))\n    axes = axes.flatten()\n    \n    ds_iter = iter(ds)\n    for n_sample in range(n_samples):\n        idx = n_sample*3\n        example_sample = next(ds_iter)\n        generated_a_sample = generator_a.predict(example_sample)\n        generated_b_sample = generator_b.predict(generated_a_sample)\n        \n        axes[idx].set_title('Input image', fontsize=18)\n        axes[idx].imshow(example_sample[0] * 0.5 + 0.5)\n        axes[idx].axis('off')\n        \n        axes[idx+1].set_title('Generated image', fontsize=18)\n        axes[idx+1].imshow(generated_a_sample[0] * 0.5 + 0.5)\n        axes[idx+1].axis('off')\n        \n        axes[idx+2].set_title('Cycled image', fontsize=18)\n        axes[idx+2].imshow(generated_b_sample[0] * 0.5 + 0.5)\n        axes[idx+2].axis('off')\n        \n    plt.show()\n\ndef create_gif(images_path, gif_path):\n    images = []\n    filenames = glob.glob(images_path)\n    filenames.sort(key=lambda x: int(''.join(filter(str.isdigit, x))))\n    for epoch, filename in enumerate(filenames):\n        img = PIL.ImageDraw.Image.open(filename)\n        ImageDraw.Draw(img).text((0, 0),  # Coordinates\n                                 f'Epoch {epoch+1}')\n        images.append(img)\n    imageio.mimsave(gif_path, images, fps=2) # Save gif\n        \ndef predict_and_save(input_ds, generator_model, output_path):\n    i = 1\n    for img in input_ds:\n        prediction = generator_model(img, training=False)[0].numpy() # make predition\n        prediction = (prediction * 127.5 + 127.5).astype(np.uint8)   # re-scale\n        im = PIL.Image.fromarray(prediction)\n        im.save(f'{output_path}{str(i)}.jpg')\n        i += 1","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Auxiliar functions (model)\n\nHere we the building blocks of our models:\n- Encoder block: Apply convolutional filters while also reducing data resolution and increasing features.\n- Decoder block: Apply convolutional filters while also increasing data resolution and decreasing features.\n- Transformer block: Apply convolutional filters to find relevant data patterns and keeps features constant.","metadata":{}},{"cell_type":"code","source":"conv_initializer = tf.random_normal_initializer(mean=0.0, stddev=0.02)\ngamma_initializer = tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n    \ndef encoder_block(input_layer, filters, size=3, strides=2, apply_instancenorm=True, activation=L.ReLU(), name='block_x'):\n    block = L.Conv2D(filters, size, \n                     strides=strides, \n                     padding='same', \n                     use_bias=False, \n                     kernel_initializer=conv_initializer, \n                     name=f'encoder_{name}')(input_layer)\n\n    if apply_instancenorm:\n        block = tfa.layers.InstanceNormalization(gamma_initializer=gamma_initializer)(block)\n        \n    block = activation(block)\n\n    return block\n\ndef transformer_block(input_layer, size=3, strides=1, name='block_x'):\n    filters = input_layer.shape[-1]\n    \n    block = L.Conv2D(filters, size, strides=strides, padding='same', use_bias=False, \n                     kernel_initializer=conv_initializer, name=f'transformer_{name}_1')(input_layer)\n#     block = tfa.layers.InstanceNormalization(gamma_initializer=gamma_initializer)(block)\n    block = L.ReLU()(block)\n    \n    block = L.Conv2D(filters, size, strides=strides, padding='same', use_bias=False, \n                     kernel_initializer=conv_initializer, name=f'transformer_{name}_2')(block)\n#     block = tfa.layers.InstanceNormalization(gamma_initializer=gamma_initializer)(block)\n    \n    block = L.Add()([block, input_layer])\n\n    return block\n\ndef decoder_block(input_layer, filters, size=3, strides=2, apply_instancenorm=True, name='block_x'):\n    block = L.Conv2DTranspose(filters, size, \n                              strides=strides, \n                              padding='same', \n                              use_bias=False, \n                              kernel_initializer=conv_initializer, \n                              name=f'decoder_{name}')(input_layer)\n\n    if apply_instancenorm:\n        block = tfa.layers.InstanceNormalization(gamma_initializer=gamma_initializer)(block)\n\n    block = L.ReLU()(block)\n    \n    return block\n\n# Resized convolution\ndef decoder_rc_block(input_layer, filters, size=3, strides=1, apply_instancenorm=True, name='block_x'):\n    block = tf.image.resize(images=input_layer, method='bilinear', \n                            size=(input_layer.shape[1]*2, input_layer.shape[2]*2))\n    \n#     block = tf.pad(block, [[0, 0], [1, 1], [1, 1], [0, 0]], \"SYMMETRIC\") # Works only with GPU\n#     block = L.Conv2D(filters, size, strides=strides, padding='valid', use_bias=False, # Works only with GPU\n    block = L.Conv2D(filters, size, \n                     strides=strides, \n                     padding='same', \n                     use_bias=False, \n                     kernel_initializer=conv_initializer, \n                     name=f'decoder_{name}')(block)\n\n    if apply_instancenorm:\n        block = tfa.layers.InstanceNormalization(gamma_initializer=gamma_initializer)(block)\n\n    block = L.ReLU()(block)\n    \n    return block","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Generator model\n\nThe `generator` is responsible for generating images from a specific domain. `CycleGAN` architecture has two generators, in this context we will have one `generator` that will take `photos` and generate `Monet paints`, and the other `generator` will take `Monet paintings` and generate `photos`.\n\nBellow, we have the architecture of the original `CycleGAN` `generator`, ours have some changes to improve performance on this task.\n\n<center><img src='https://github.com/dimitreOliveira/MachineLearning/blob/master/Kaggle/I%E2%80%99m%20Something%20of%20a%20Painter%20Myself/generator_architecture.png?raw=true' height=250></center>","metadata":{}},{"cell_type":"code","source":"def generator_fn(height=HEIGHT, width=WIDTH, channels=CHANNELS, transformer_blocks=TRANSFORMER_BLOCKS):\n    OUTPUT_CHANNELS = 3\n    inputs = L.Input(shape=[height, width, channels], name='input_image')\n\n    # Encoder\n    enc_1 = encoder_block(inputs, 64,  7, 1, apply_instancenorm=False, activation=L.ReLU(), name='block_1') # (bs, 256, 256, 64)\n    enc_2 = encoder_block(enc_1, 128, 3, 2, apply_instancenorm=True, activation=L.ReLU(), name='block_2')   # (bs, 128, 128, 128)\n    enc_3 = encoder_block(enc_2, 256, 3, 2, apply_instancenorm=True, activation=L.ReLU(), name='block_3')   # (bs, 64, 64, 256)\n    \n    # Transformer\n    x = enc_3\n    for n in range(transformer_blocks):\n        x = transformer_block(x, 3, 1, name=f'block_{n+1}') # (bs, 64, 64, 256)\n\n    # Decoder\n    x_skip = L.Concatenate(name='enc_dec_skip_1')([x, enc_3]) # encoder - decoder skip connection\n    \n    dec_1 = decoder_block(x_skip, 128, 3, 2, apply_instancenorm=True, name='block_1') # (bs, 128, 128, 128)\n    x_skip = L.Concatenate(name='enc_dec_skip_2')([dec_1, enc_2]) # encoder - decoder skip connection\n    \n    dec_2 = decoder_block(x_skip, 64,  3, 2, apply_instancenorm=True, name='block_2') # (bs, 256, 256, 64)\n    x_skip = L.Concatenate(name='enc_dec_skip_3')([dec_2, enc_1]) # encoder - decoder skip connection\n\n    outputs = last = L.Conv2D(OUTPUT_CHANNELS, 7, \n                              strides=1, padding='same', \n                              kernel_initializer=conv_initializer, \n                              use_bias=False, \n                              activation='tanh', \n                              name='decoder_output_block')(x_skip) # (bs, 256, 256, 3)\n\n    generator = Model(inputs, outputs)\n    \n    return generator\n\nsample_generator = generator_fn()\nsample_generator.summary()","metadata":{"_kg_hide-input":false,"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Discriminator model\n\n\nThe `discriminator` is responsible for differentiating real images from images that have been generated by a `generator` model.\n\nBellow, we have the architecture of the original `CycleGAN` `discriminator`, again, ours have some changes to improve performance on this task.\n\n<center><img src='https://github.com/dimitreOliveira/MachineLearning/blob/master/Kaggle/I%E2%80%99m%20Something%20of%20a%20Painter%20Myself/discriminator_architecture.png?raw=true' height=550, width=550></center>","metadata":{}},{"cell_type":"code","source":"def discriminator_fn(height=HEIGHT, width=WIDTH, channels=CHANNELS):\n    inputs = L.Input(shape=[height, width, channels], name='input_image')\n    #inputs_patch = L.experimental.preprocessing.RandomCrop(height=70, width=70, name='input_image_patch')(inputs) # Works only with GPU\n\n    # Encoder    \n    x = encoder_block(inputs, 64,  4, 2, apply_instancenorm=False, activation=L.LeakyReLU(0.2), name='block_1') # (bs, 128, 128, 64)\n    x = encoder_block(x, 128, 4, 2, apply_instancenorm=True, activation=L.LeakyReLU(0.2), name='block_2')       # (bs, 64, 64, 128)\n    x = encoder_block(x, 256, 4, 2, apply_instancenorm=True, activation=L.LeakyReLU(0.2), name='block_3')       # (bs, 32, 32, 256)\n    x = encoder_block(x, 512, 4, 1, apply_instancenorm=True, activation=L.LeakyReLU(0.2), name='block_4')       # (bs, 32, 32, 512)\n\n    outputs = L.Conv2D(1, 4, strides=1, padding='valid', kernel_initializer=conv_initializer)(x)                # (bs, 29, 29, 1)\n    \n    discriminator = Model(inputs, outputs)\n    \n    return discriminator\n\n\nsample_discriminator = discriminator_fn()\nsample_discriminator.summary()","metadata":{"_kg_hide-input":false,"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Build model (CycleGAN)","metadata":{}},{"cell_type":"code","source":"with strategy.scope():\n    monet_generator = generator_fn(height=None, width=None, transformer_blocks=TRANSFORMER_BLOCKS) # transforms photos to Monet-esque paintings\n    photo_generator = generator_fn(height=None, width=None, transformer_blocks=TRANSFORMER_BLOCKS) # transforms Monet paintings to be more like photos\n\n    monet_discriminator = discriminator_fn(height=None, width=None) # differentiates real Monet paintings and generated Monet paintings\n    photo_discriminator = discriminator_fn(height=None, width=None) # differentiates real photos and generated photos\n\n\nclass CycleGan(Model):\n    def __init__(\n        self,\n        monet_generator,\n        photo_generator,\n        monet_discriminator,\n        photo_discriminator,\n        lambda_cycle=10,\n    ):\n        super(CycleGan, self).__init__()\n        self.m_gen = monet_generator\n        self.p_gen = photo_generator\n        self.m_disc = monet_discriminator\n        self.p_disc = photo_discriminator\n        self.lambda_cycle = lambda_cycle\n        \n    def compile(\n        self,\n        m_gen_optimizer,\n        p_gen_optimizer,\n        m_disc_optimizer,\n        p_disc_optimizer,\n        gen_loss_fn,\n        disc_loss_fn,\n        cycle_loss_fn,\n        identity_loss_fn\n    ):\n        super(CycleGan, self).compile()\n        self.m_gen_optimizer = m_gen_optimizer\n        self.p_gen_optimizer = p_gen_optimizer\n        self.m_disc_optimizer = m_disc_optimizer\n        self.p_disc_optimizer = p_disc_optimizer\n        self.gen_loss_fn = gen_loss_fn\n        self.disc_loss_fn = disc_loss_fn\n        self.cycle_loss_fn = cycle_loss_fn\n        self.identity_loss_fn = identity_loss_fn\n        \n    def train_step(self, batch_data):\n        real_monet, real_photo = batch_data\n        \n        with tf.GradientTape(persistent=True) as tape:\n            # photo to monet back to photo\n            fake_monet = self.m_gen(real_photo, training=True)\n            cycled_photo = self.p_gen(fake_monet, training=True)\n\n            # monet to photo back to monet\n            fake_photo = self.p_gen(real_monet, training=True)\n            cycled_monet = self.m_gen(fake_photo, training=True)\n\n            # generating itself\n            same_monet = self.m_gen(real_monet, training=True)\n            same_photo = self.p_gen(real_photo, training=True)\n\n            # discriminator used to check, inputing real images\n            disc_real_monet = self.m_disc(real_monet, training=True)\n            disc_real_photo = self.p_disc(real_photo, training=True)\n\n            # discriminator used to check, inputing fake images\n            disc_fake_monet = self.m_disc(fake_monet, training=True)\n            disc_fake_photo = self.p_disc(fake_photo, training=True)\n\n            # evaluates generator loss\n            monet_gen_loss = self.gen_loss_fn(disc_fake_monet)\n            photo_gen_loss = self.gen_loss_fn(disc_fake_photo)\n\n            # evaluates total cycle consistency loss\n            total_cycle_loss = self.cycle_loss_fn(real_monet, cycled_monet, self.lambda_cycle) + self.cycle_loss_fn(real_photo, cycled_photo, self.lambda_cycle)\n\n            # evaluates total generator loss\n            total_monet_gen_loss = monet_gen_loss + total_cycle_loss + self.identity_loss_fn(real_monet, same_monet, self.lambda_cycle)\n            total_photo_gen_loss = photo_gen_loss + total_cycle_loss + self.identity_loss_fn(real_photo, same_photo, self.lambda_cycle)\n\n            # evaluates discriminator loss\n            monet_disc_loss = self.disc_loss_fn(disc_real_monet, disc_fake_monet)\n            photo_disc_loss = self.disc_loss_fn(disc_real_photo, disc_fake_photo)\n\n        # Calculate the gradients for generator and discriminator\n        monet_generator_gradients = tape.gradient(total_monet_gen_loss,\n                                                  self.m_gen.trainable_variables)\n        photo_generator_gradients = tape.gradient(total_photo_gen_loss,\n                                                  self.p_gen.trainable_variables)\n\n        monet_discriminator_gradients = tape.gradient(monet_disc_loss,\n                                                      self.m_disc.trainable_variables)\n        photo_discriminator_gradients = tape.gradient(photo_disc_loss,\n                                                      self.p_disc.trainable_variables)\n\n        # Apply the gradients to the optimizer\n        self.m_gen_optimizer.apply_gradients(zip(monet_generator_gradients,\n                                                 self.m_gen.trainable_variables))\n\n        self.p_gen_optimizer.apply_gradients(zip(photo_generator_gradients,\n                                                 self.p_gen.trainable_variables))\n\n        self.m_disc_optimizer.apply_gradients(zip(monet_discriminator_gradients,\n                                                  self.m_disc.trainable_variables))\n\n        self.p_disc_optimizer.apply_gradients(zip(photo_discriminator_gradients,\n                                                  self.p_disc.trainable_variables))\n        \n        return {'monet_gen_loss': total_monet_gen_loss,\n                'photo_gen_loss': total_photo_gen_loss,\n                'monet_disc_loss': monet_disc_loss,\n                'photo_disc_loss': photo_disc_loss\n               }","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Loss functions","metadata":{}},{"cell_type":"code","source":"with strategy.scope():\n    # Discriminator loss {0: fake, 1: real} (The discriminator loss outputs the average of the real and generated loss)\n    def discriminator_loss(real, generated):\n        real_loss = losses.BinaryCrossentropy(from_logits=True, reduction=losses.Reduction.NONE)(tf.ones_like(real), real)\n\n        generated_loss = losses.BinaryCrossentropy(from_logits=True, reduction=losses.Reduction.NONE)(tf.zeros_like(generated), generated)\n\n        total_disc_loss = real_loss + generated_loss\n\n        return total_disc_loss * 0.5\n    \n    # Generator loss\n    def generator_loss(generated):\n        return losses.BinaryCrossentropy(from_logits=True, reduction=losses.Reduction.NONE)(tf.ones_like(generated), generated)\n    \n    \n    # Cycle consistency loss (measures if original photo and the twice transformed photo to be similar to one another)\n    with strategy.scope():\n        def calc_cycle_loss(real_image, cycled_image, LAMBDA):\n            loss1 = tf.reduce_mean(tf.abs(real_image - cycled_image))\n\n            return LAMBDA * loss1\n\n    # Identity loss (compares the image with its generator (i.e. photo with photo generator))\n    with strategy.scope():\n        def identity_loss(real_image, same_image, LAMBDA):\n            loss = tf.reduce_mean(tf.abs(real_image - same_image))\n            return LAMBDA * 0.5 * loss","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Learning rate schedule\n\nThe original `CycleGAN` implementation used a `constant learning rate schedule with a linear decay`, I also found that the linear decay phase seems to be good at making the model more stable at the last epochs, you can check how the `generator` changes in a more conservative rate by the end looking at the `gif` images by the end.","metadata":{}},{"cell_type":"code","source":"@tf.function\ndef linear_schedule_with_warmup(step):\n    \"\"\" Create a schedule with a learning rate that decreases linearly after\n    linearly increasing during a warmup period.\n    \"\"\"\n    lr_start   = 2e-4\n    lr_max     = 2e-4\n    lr_min     = 0.\n    \n    steps_per_epoch = int(max(n_monet_samples, n_photo_samples)//BATCH_SIZE)\n    total_steps = EPOCHS * steps_per_epoch\n    warmup_steps = 1\n    hold_max_steps = total_steps * 0.8\n    \n    if step < warmup_steps:\n        lr = (lr_max - lr_start) / warmup_steps * step + lr_start\n    elif step < warmup_steps + hold_max_steps:\n        lr = lr_max\n    else:\n        lr = lr_max * ((total_steps - step) / (total_steps - warmup_steps - hold_max_steps))\n        if lr_min is not None:\n            lr = tf.math.maximum(lr_min, lr)\n\n    return lr\n\nsteps_per_epoch = int(max(n_monet_samples, n_photo_samples)//BATCH_SIZE)\ntotal_steps = EPOCHS * steps_per_epoch\nrng = [i for i in range(0, total_steps, 50)]\ny = [linear_schedule_with_warmup(x) for x in rng]\n\nsns.set(style=\"whitegrid\")\nfig, ax = plt.subplots(figsize=(20, 6))\nplt.plot(rng, y)\nprint(f'{EPOCHS} total epochs and {steps_per_epoch} steps per epoch')\nprint(f'Learning rate schedule: {y[0]:.3g} to {max(y):.3g} to {y[-1]:.3g}')","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train","metadata":{}},{"cell_type":"code","source":"with strategy.scope():\n    # Create generators\n    lr_monet_gen = lambda: linear_schedule_with_warmup(tf.cast(monet_generator_optimizer.iterations, tf.float32))\n    lr_photo_gen = lambda: linear_schedule_with_warmup(tf.cast(photo_generator_optimizer.iterations, tf.float32))\n    \n    monet_generator_optimizer = optimizers.Adam(learning_rate=lr_monet_gen, beta_1=0.5)\n    photo_generator_optimizer = optimizers.Adam(learning_rate=lr_photo_gen, beta_1=0.5)\n\n    # Create discriminators\n    lr_monet_disc = lambda: linear_schedule_with_warmup(tf.cast(monet_discriminator_optimizer.iterations, tf.float32))\n    lr_photo_disc = lambda: linear_schedule_with_warmup(tf.cast(photo_discriminator_optimizer.iterations, tf.float32))\n    \n    monet_discriminator_optimizer = optimizers.Adam(learning_rate=lr_monet_disc, beta_1=0.5)\n    photo_discriminator_optimizer = optimizers.Adam(learning_rate=lr_photo_disc, beta_1=0.5)\n\n    \n    # Create GAN\n    gan_model = CycleGan(monet_generator, photo_generator, \n                         monet_discriminator, photo_discriminator)\n\n    gan_model.compile(m_gen_optimizer=monet_generator_optimizer,\n                      p_gen_optimizer=photo_generator_optimizer,\n                      m_disc_optimizer=monet_discriminator_optimizer,\n                      p_disc_optimizer=photo_discriminator_optimizer,\n                      gen_loss_fn=generator_loss,\n                      disc_loss_fn=discriminator_loss,\n                      cycle_loss_fn=calc_cycle_loss,\n                      identity_loss_fn=identity_loss)","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Create dataset\nmonet_ds = get_dataset(MONET_FILENAMES, augment=data_augment, batch_size=BATCH_SIZE)\nphoto_ds = get_dataset(PHOTO_FILENAMES, augment=data_augment, batch_size=BATCH_SIZE)\ngan_ds = tf.data.Dataset.zip((monet_ds, photo_ds))\n\nphoto_ds_eval = get_dataset(PHOTO_FILENAMES, repeat=False, shuffle=False, batch_size=1)\nmonet_ds_eval = get_dataset(MONET_FILENAMES, repeat=False, shuffle=False, batch_size=1)\n\n# Callbacks\nclass GANMonitor(Callback):\n    \"\"\"A callback to generate and save images after each epoch\"\"\"\n\n    def __init__(self, num_img=1, monet_path='monet', photo_path='photo'):\n        self.num_img = num_img\n        self.monet_path = monet_path\n        self.photo_path = photo_path\n        # Create directories to save the generate images\n        if not os.path.exists(self.monet_path):\n            os.makedirs(self.monet_path)\n        if not os.path.exists(self.photo_path):\n            os.makedirs(self.photo_path)\n\n    def on_epoch_end(self, epoch, logs=None):\n        # Monet generated images\n        for i, img in enumerate(photo_ds_eval.take(self.num_img)):\n            prediction = monet_generator(img, training=False)[0].numpy()\n            prediction = (prediction * 127.5 + 127.5).astype(np.uint8)\n            prediction = PIL.Image.fromarray(prediction)\n            prediction.save(f'{self.monet_path}/generated_{i}_{epoch+1}.png')\n            \n        # Photo generated images\n        for i, img in enumerate(monet_ds_eval.take(self.num_img)):\n            prediction = photo_generator(img, training=False)[0].numpy()\n            prediction = (prediction * 127.5 + 127.5).astype(np.uint8)\n            prediction = PIL.Image.fromarray(prediction)\n            prediction.save(f'{self.photo_path}/generated_{i}_{epoch+1}.png')","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"history = gan_model.fit(gan_ds, \n                        epochs=EPOCHS, \n                        callbacks=[GANMonitor()], \n                        steps_per_epoch=(max(n_monet_samples, n_photo_samples)//BATCH_SIZE), \n                        verbose=2).history","metadata":{"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We can see the generators progress at each epoch by creating a `gif` that is a generated image at each epoch.\n\n## Monet generation GIF","metadata":{}},{"cell_type":"code","source":"# Got the idea from https://www.kaggle.com/matkneky/monet-cyclegan-trials\n# Create GIFs\ncreate_gif('/kaggle/working/monet/*.png', 'monet.gif') # Create monet gif","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<img src='monet.gif' width=350>","metadata":{}},{"cell_type":"markdown","source":"## Photo generation GIF","metadata":{}},{"cell_type":"code","source":"create_gif('/kaggle/working/photo/*.png', 'photo.gif') # Create photo gif","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<img src='photo.gif' width=350>","metadata":{}},{"cell_type":"markdown","source":"# Evaluating generator models\n\nHere we are going to evaluate the generator models including how good is the generator cycle, this means that we will get a photo to generate a Monet picture from it, then use the generated picture to generate the original photo.\n\n## Photo (input) -> Monet (generated) -> Photo (generated)","metadata":{}},{"cell_type":"code","source":"evaluate_cycle(photo_ds_eval.take(2), monet_generator, photo_generator, n_samples=2)","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Here we will do the same process but starting with a Monet picture.\n\n## Monet (input) -> Photo (generated) -> Monet (generated)","metadata":{}},{"cell_type":"code","source":"evaluate_cycle(monet_ds_eval.take(2), photo_generator, monet_generator, n_samples=2)","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Visualize predictions\n\nA common issue with images generated by GANs is that the often show some undisered artifacts, a very common on is known as \"[checkerboard artifacts](https://distill.pub/2016/deconv-checkerboard/)\", a good practice is to inspect some of the images to see its quality and if some of these undisered artifacts are present.","metadata":{}},{"cell_type":"code","source":"display_generated_samples(photo_ds_eval.take(8), monet_generator, 8)","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Make predictions","metadata":{}},{"cell_type":"code","source":"%%time\nos.makedirs('../images/') # Create folder to save generated images\npredict_and_save(photo_ds_eval, monet_generator, '../images/')","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Submission file","metadata":{}},{"cell_type":"code","source":"shutil.make_archive('/kaggle/working/images/', 'zip', '../images')\n\nprint(f\"Generated samples: {len([name for name in os.listdir('../images/') if os.path.isfile(os.path.join('../images/', name))])}\")","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Output models","metadata":{}},{"cell_type":"code","source":"monet_generator.save('monet_generator.h5')\nphoto_generator.save('photo_generator.h5')\nmonet_discriminator.save('monet_discriminator.h5')\nphoto_discriminator.save('photo_discriminator.h5')","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]}]}