{"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":"code","source":"#Importing dependencies\n\nimport os, random, json, PIL, shutil, re, imageio, glob\nos.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\nos.environ[\"AUTOGRAPH_VERBOSITY\"] = \"0\"\n\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)\n\n#Initializing a TPU to do computations with.\ntry:\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n    print('Device:', tpu.master())\n    tf.config.experimental_connect_to_cluster(tpu)\n    tf.tpu.experimental.initialize_tpu_system(tpu)\n    strategy = tf.distribute.experimental.TPUStrategy(tpu)\nexcept:\n    strategy = tf.distribute.get_strategy()\nprint('Number of replicas:', strategy.num_replicas_in_sync)\n\nAUTOTUNE = tf.data.experimental.AUTOTUNE\n    \nprint(tf.__version__)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-05-13T08:16:09.645377Z","iopub.execute_input":"2022-05-13T08:16:09.646713Z","iopub.status.idle":"2022-05-13T08:16:20.604239Z","shell.execute_reply.started":"2022-05-13T08:16:09.646648Z","shell.execute_reply":"2022-05-13T08:16:20.603410Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#Getting the GCS path and loading up the files.\nGCS_PATH = KaggleDatasets().get_gcs_path()\nmonet_files = tf.io.gfile.glob(str(GCS_PATH + '/monet_tfrec/*.tfrec'))\nphoto_files = tf.io.gfile.glob(str(GCS_PATH + '/photo_tfrec/*.tfrec'))","metadata":{"execution":{"iopub.status.busy":"2022-05-13T08:16:20.606371Z","iopub.execute_input":"2022-05-13T08:16:20.606906Z","iopub.status.idle":"2022-05-13T08:16:21.205561Z","shell.execute_reply.started":"2022-05-13T08:16:20.606859Z","shell.execute_reply":"2022-05-13T08:16:21.204550Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"These are the model parameters I used in this case. These can be changed for different GANs.","metadata":{}},{"cell_type":"code","source":"#Model Parameters \n\nHEIGHT = 256\nWIDTH = 256\nHEIGHT_RESIZE = 128\nWIDTH_RESIZE = 128\nCHANNELS = 3\nBATCH_SIZE = 12\nEPOCHS = 25\nTRANSFORMER_BLOCKS = 6\nGENERATOR_LR = 2e-4\nDISCRIMINATOR_LR = 2e-4","metadata":{"execution":{"iopub.status.busy":"2022-05-13T08:16:21.207308Z","iopub.execute_input":"2022-05-13T08:16:21.207703Z","iopub.status.idle":"2022-05-13T08:16:21.213994Z","shell.execute_reply.started":"2022-05-13T08:16:21.207651Z","shell.execute_reply":"2022-05-13T08:16:21.213205Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"I will now create some helper functions to be used later throughout the data preprocessing, training, and visualization of results since these are images.","metadata":{}},{"cell_type":"code","source":"image_size = [256,256]\n\n#Normalization function\ndef 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\n#Define a function to read the tfrec file extensions into images.\ndef read_tfrec(file):\n    tfrecord_format = {\n          \"image_name\": tf.io.FixedLenFeature([], tf.string),\n          \"image\": tf.io.FixedLenFeature([], tf.string),\n          \"target\": tf.io.FixedLenFeature([], tf.string)\n    }  #Setting up the file format.\n\n    #decode the image \n    example = tf.io.parse_single_example(file, tfrecord_format)\n    image = example['image'] \n    image = tf.image.decode_jpeg(image, channels = 3) # used 3 channels because it's in RGB\n    image = (tf.cast(image, tf.float32) / 127.5) - 1\n    image = tf.reshape(image, [*image_size, 3])\n\n    return image\n\n#Load up the images into a dataset\ndef load_images(filenames, labeled=True, ordered=False):\n    dataset = tf.data.TFRecordDataset(filenames)\n    dataset = dataset.map(read_tfrec, num_parallel_calls=AUTOTUNE)\n    return dataset\n\ndef get_dataset(filenames,repeat=True, shuffle=True, batch_size=1):\n    dataset = load_images(filenames)\n    \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":{"execution":{"iopub.status.busy":"2022-05-13T08:16:21.216045Z","iopub.execute_input":"2022-05-13T08:16:21.216293Z","iopub.status.idle":"2022-05-13T08:16:21.253426Z","shell.execute_reply.started":"2022-05-13T08:16:21.216254Z","shell.execute_reply":"2022-05-13T08:16:21.252128Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#Getting a separate dataset for the monet photos and the original ones.\nmonet_ds = load_images(monet_files, labeled=True).batch(1)\nphoto_ds = load_images(photo_files, labeled=True).batch(1)","metadata":{"execution":{"iopub.status.busy":"2022-05-13T08:16:21.254796Z","iopub.execute_input":"2022-05-13T08:16:21.255240Z","iopub.status.idle":"2022-05-13T08:16:21.327373Z","shell.execute_reply.started":"2022-05-13T08:16:21.255195Z","shell.execute_reply":"2022-05-13T08:16:21.326349Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#Getting examples from each to take a look at them.\n\nexample_monet = next(iter(monet_ds))\nexample_photo = next(iter(photo_ds))\n\nplt.subplot(121)\nplt.title('Photo')\nplt.imshow(example_photo[0] * 0.5 + 0.5)\n\nplt.subplot(122)\nplt.title('Monet')\nplt.imshow(example_monet[0] * 0.5 + 0.5)","metadata":{"execution":{"iopub.status.busy":"2022-05-13T08:16:21.328935Z","iopub.execute_input":"2022-05-13T08:16:21.329263Z","iopub.status.idle":"2022-05-13T08:16:21.985473Z","shell.execute_reply.started":"2022-05-13T08:16:21.329223Z","shell.execute_reply":"2022-05-13T08:16:21.984317Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"I will now create the methods to be implemented in the CycleGAN class I will create later to avoid clustering. These methods are:\n\n1-Encoder\n\n2-Decoder\n\n3-Transformer\n\nBoth of these will be used in the generator and discrimnator function\n\n5-Generator\n\n6-Discriminator\n\nLosses: \n\n7-Generator loss\n\n8-Discriminator loss\n\n9-Cycle consistency loss\n\n10-Identity loss\n","metadata":{}},{"cell_type":"code","source":"conv_initializer = tf.random_normal_initializer(mean=0.0, stddev=0.04)\ngamma_initializer = tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.04)\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 = 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    \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#Resize decoder too! \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    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":{"execution":{"iopub.status.busy":"2022-05-13T08:16:21.987304Z","iopub.execute_input":"2022-05-13T08:16:21.987666Z","iopub.status.idle":"2022-05-13T08:16:22.013134Z","shell.execute_reply.started":"2022-05-13T08:16:21.987624Z","shell.execute_reply":"2022-05-13T08:16:22.012292Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"I will create general generators and discriminators since they are both used twice in CycleGAN (One for the original photo and one for the \"styled\" photo).\n\n\nSkip connections are used here to get around the vanishing gradient problem.","metadata":{}},{"cell_type":"code","source":"def gen_func(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\ndef disc_func(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","metadata":{"execution":{"iopub.status.busy":"2022-05-13T08:16:22.014744Z","iopub.execute_input":"2022-05-13T08:16:22.015296Z","iopub.status.idle":"2022-05-13T08:16:22.034939Z","shell.execute_reply.started":"2022-05-13T08:16:22.015258Z","shell.execute_reply":"2022-05-13T08:16:22.033826Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#Initialize both generators and both discriminators\nwith strategy.scope():\n    monet_gen = gen_func(height=None, width=None, transformer_blocks=TRANSFORMER_BLOCKS) # transforms photos to Monet-esque paintings\n    photo_gen = gen_func(height=None, width=None, transformer_blocks=TRANSFORMER_BLOCKS) # transforms Monet paintings to be more like photos\n\n    monet_disc = disc_func(height=None, width=None) # differentiates real Monet paintings and generated Monet paintings\n    photo_disc = disc_func(height=None, width=None) # differentiates real photos and generated photos","metadata":{"execution":{"iopub.status.busy":"2022-05-13T08:16:22.036726Z","iopub.execute_input":"2022-05-13T08:16:22.037442Z","iopub.status.idle":"2022-05-13T08:16:24.125668Z","shell.execute_reply.started":"2022-05-13T08:16:22.037384Z","shell.execute_reply":"2022-05-13T08:16:24.124719Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Implementing the 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, gen):\n        real_loss = losses.BinaryCrossentropy(from_logits=True, reduction=losses.Reduction.NONE)(tf.ones_like(real), real)\n\n        gen_loss = losses.BinaryCrossentropy(from_logits=True, reduction=losses.Reduction.NONE)(tf.zeros_like(gen), gen)\n\n        total_disc_loss = real_loss + gen_loss\n\n        return total_disc_loss * 0.5\n    \n    # Generator loss\n    def generator_loss(gen):\n        return losses.BinaryCrossentropy(from_logits=True, reduction=losses.Reduction.NONE)(tf.ones_like(gen), gen)\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":{"execution":{"iopub.status.busy":"2022-05-13T08:16:24.129139Z","iopub.execute_input":"2022-05-13T08:16:24.129545Z","iopub.status.idle":"2022-05-13T08:16:24.141885Z","shell.execute_reply.started":"2022-05-13T08:16:24.129506Z","shell.execute_reply":"2022-05-13T08:16:24.140771Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CycleGan(Model):\n    def __init__(\n        self,\n        monet_gen,\n        photo_gen,\n        monet_disc,\n        photo_disc,\n        lambda_cycle=10,\n    ):\n        super(CycleGan, self).__init__()\n        self.m_gen = monet_gen\n        self.p_gen = photo_gen\n        self.m_disc = monet_disc\n        self.p_disc = photo_disc\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            \n        # Calculate the gradients for generator and discriminator\n        monet_gen_gradients = tape.gradient(total_monet_gen_loss,\n                                                  self.m_gen.trainable_variables)\n        photo_gen_gradients = tape.gradient(total_photo_gen_loss,\n                                                  self.p_gen.trainable_variables)\n\n        monet_disc_gradients = tape.gradient(monet_disc_loss,\n                                                      self.m_disc.trainable_variables)\n        photo_disc_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_gen_gradients,\n                                                 self.m_gen.trainable_variables))\n\n        self.p_gen_optimizer.apply_gradients(zip(photo_gen_gradients,\n                                                 self.p_gen.trainable_variables))\n\n        self.m_disc_optimizer.apply_gradients(zip(monet_disc_gradients,\n                                                  self.m_disc.trainable_variables))\n\n        self.p_disc_optimizer.apply_gradients(zip(photo_disc_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":{"execution":{"iopub.status.busy":"2022-05-13T08:16:24.143845Z","iopub.execute_input":"2022-05-13T08:16:24.144299Z","iopub.status.idle":"2022-05-13T08:16:24.171411Z","shell.execute_reply.started":"2022-05-13T08:16:24.144265Z","shell.execute_reply":"2022-05-13T08:16:24.169852Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"I have to do a learning rate schedule to decrease at some point to avoid diverging when being close to the local minima. (This function is copied from a collaborator).","metadata":{}},{"cell_type":"code","source":"def 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_files)\nn_photo_samples = count_data_items(photo_files)\n\n@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":{"execution":{"iopub.status.busy":"2022-05-13T08:16:24.173489Z","iopub.execute_input":"2022-05-13T08:16:24.173876Z","iopub.status.idle":"2022-05-13T08:16:27.502622Z","shell.execute_reply.started":"2022-05-13T08:16:24.173837Z","shell.execute_reply":"2022-05-13T08:16:27.501889Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Now it's time to train our model!","metadata":{}},{"cell_type":"code","source":"with strategy.scope():\n    # Create generators\n    lr_monet_gen = lambda: linear_schedule_with_warmup(tf.cast(monet_gen_optimizer.iterations, tf.float32))\n    lr_photo_gen = lambda: linear_schedule_with_warmup(tf.cast(photo_gen_optimizer.iterations, tf.float32))\n    \n    monet_gen_optimizer = optimizers.Adam(learning_rate=lr_monet_gen, beta_1=0.5)\n    photo_gen_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_disc_optimizer.iterations, tf.float32))\n    lr_photo_disc = lambda: linear_schedule_with_warmup(tf.cast(photo_disc_optimizer.iterations, tf.float32))\n    \n    monet_disc_optimizer = optimizers.Adam(learning_rate=lr_monet_disc, beta_1=0.5)\n    photo_disc_optimizer = optimizers.Adam(learning_rate=lr_photo_disc, beta_1=0.5)\n\n    \n    # Create GAN\n    gan_model = CycleGan(monet_gen, photo_gen, \n                         monet_disc, photo_disc)\n\n    gan_model.compile(m_gen_optimizer=monet_gen_optimizer,\n                      p_gen_optimizer=photo_gen_optimizer,\n                      m_disc_optimizer=monet_disc_optimizer,\n                      p_disc_optimizer=photo_disc_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":{"execution":{"iopub.status.busy":"2022-05-13T08:16:27.503930Z","iopub.execute_input":"2022-05-13T08:16:27.504694Z","iopub.status.idle":"2022-05-13T08:16:27.573806Z","shell.execute_reply.started":"2022-05-13T08:16:27.504645Z","shell.execute_reply":"2022-05-13T08:16:27.572879Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The callbacks part here is to construct the gif transforming from one style to another. (The code for it is not written by me).","metadata":{}},{"cell_type":"code","source":"# Create dataset\nmonet_ds = get_dataset(monet_files, batch_size=BATCH_SIZE)\nphoto_ds = get_dataset(photo_files, batch_size=BATCH_SIZE)\ngan_ds = tf.data.Dataset.zip((monet_ds, photo_ds))\n\nphoto_ds_eval = get_dataset(photo_files, repeat=False, shuffle=False, batch_size=1)\nmonet_ds_eval = get_dataset(monet_files, 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_gen(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_gen(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":{"execution":{"iopub.status.busy":"2022-05-13T08:16:27.575267Z","iopub.execute_input":"2022-05-13T08:16:27.575527Z","iopub.status.idle":"2022-05-13T08:16:27.704575Z","shell.execute_reply.started":"2022-05-13T08:16:27.575499Z","shell.execute_reply":"2022-05-13T08:16:27.703416Z"},"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":{"execution":{"iopub.status.busy":"2022-05-13T08:16:27.705929Z","iopub.execute_input":"2022-05-13T08:16:27.706171Z","iopub.status.idle":"2022-05-13T09:21:37.535967Z","shell.execute_reply.started":"2022-05-13T08:16:27.706144Z","shell.execute_reply":"2022-05-13T09:21:37.533461Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#Displaying what's happening over all the epochs.\ncreate_gif('/kaggle/working/monet/*.png', 'monet.gif')","metadata":{"execution":{"iopub.status.busy":"2022-05-13T09:21:37.539784Z","iopub.execute_input":"2022-05-13T09:21:37.540215Z","iopub.status.idle":"2022-05-13T09:21:41.897635Z","shell.execute_reply.started":"2022-05-13T09:21:37.540162Z","shell.execute_reply":"2022-05-13T09:21:41.896705Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"create_gif('/kaggle/working/photo/*.png', 'photo.gif')","metadata":{"execution":{"iopub.status.busy":"2022-05-13T09:21:41.898869Z","iopub.execute_input":"2022-05-13T09:21:41.899113Z","iopub.status.idle":"2022-05-13T09:21:46.041290Z","shell.execute_reply.started":"2022-05-13T09:21:41.899085Z","shell.execute_reply":"2022-05-13T09:21:46.040327Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"evaluate_cycle(photo_ds_eval.take(2), monet_gen, photo_gen, n_samples=2)","metadata":{"execution":{"iopub.status.busy":"2022-05-13T09:21:46.042865Z","iopub.execute_input":"2022-05-13T09:21:46.043301Z","iopub.status.idle":"2022-05-13T09:23:30.526476Z","shell.execute_reply.started":"2022-05-13T09:21:46.043254Z","shell.execute_reply":"2022-05-13T09:23:30.525476Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"evaluate_cycle(monet_ds_eval.take(2), photo_gen, monet_gen, n_samples=2)","metadata":{"execution":{"iopub.status.busy":"2022-05-13T09:23:30.528073Z","iopub.execute_input":"2022-05-13T09:23:30.528323Z","iopub.status.idle":"2022-05-13T09:23:33.659545Z","shell.execute_reply.started":"2022-05-13T09:23:30.528293Z","shell.execute_reply":"2022-05-13T09:23:33.658565Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"display_generated_samples(photo_ds_eval.take(8), monet_gen, 8)","metadata":{"execution":{"iopub.status.busy":"2022-05-13T09:23:33.661179Z","iopub.execute_input":"2022-05-13T09:23:33.661745Z","iopub.status.idle":"2022-05-13T09:23:40.910523Z","shell.execute_reply.started":"2022-05-13T09:23:33.661702Z","shell.execute_reply":"2022-05-13T09:23:40.909520Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":" ","metadata":{},"execution_count":null,"outputs":[]}]}