{"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":"import tensorflow as tf\nfrom tensorflow import keras\nfrom tensorflow.keras import layers\nimport tensorflow_addons as tfa\n\nfrom kaggle_datasets import KaggleDatasets\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport os, random, json, PIL, shutil, re, imageio, glob\nfrom PIL import ImageDraw\nfrom IPython.display import Image\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":"2023-03-12T16:41:40.056052Z","iopub.execute_input":"2023-03-12T16:41:40.056528Z","iopub.status.idle":"2023-03-12T16:41:45.865764Z","shell.execute_reply.started":"2023-03-12T16:41:40.056436Z","shell.execute_reply":"2023-03-12T16:41:45.864723Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Load in the data\n\nWe want to keep our photo dataset and our Monet dataset separate. First, load in the filenames of the TFRecords.","metadata":{}},{"cell_type":"code","source":"GCS_PATH = KaggleDatasets().get_gcs_path()\nprint(GCS_PATH)\nMONET_FILENAMES = tf.io.gfile.glob(str(GCS_PATH + '/monet_tfrec/*.tfrec'))\nprint('Monet TFRecord Files:', len(MONET_FILENAMES))\n\nPHOTO_FILENAMES = tf.io.gfile.glob(str(GCS_PATH + '/photo_tfrec/*.tfrec'))\nprint('Photo TFRecord Files:', len(PHOTO_FILENAMES))\n\n","metadata":{"execution":{"iopub.status.busy":"2023-03-12T16:41:45.867720Z","iopub.execute_input":"2023-03-12T16:41:45.868605Z","iopub.status.idle":"2023-03-12T16:41:46.282787Z","shell.execute_reply.started":"2023-03-12T16:41:45.868567Z","shell.execute_reply":"2023-03-12T16:41:46.281753Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"All the images for the competition are already sized to 256x256. As these images are RGB images, set the channel to 3. Additionally, we need to scale the images to a [-1, 1] scale. Because we are building a generative model, we don't need the labels or the image id so we'll only return the image from the TFRecord.\n","metadata":{}},{"cell_type":"code","source":"IMAGE_SIZE = [256, 256]\n\ndef decode_image(image):\n    image = tf.image.decode_jpeg(image, channels=3)\n    image = (tf.cast(image, tf.float32) / 127.5) - 1\n    image = tf.reshape(image, [*IMAGE_SIZE, 3])\n    return image\n\ndef read_tfrecord(example):\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    }\n    example = tf.io.parse_single_example(example, tfrecord_format)\n    image = decode_image(example['image'])\n    return image","metadata":{"execution":{"iopub.status.busy":"2023-03-12T16:41:46.284166Z","iopub.execute_input":"2023-03-12T16:41:46.284726Z","iopub.status.idle":"2023-03-12T16:41:46.292058Z","shell.execute_reply.started":"2023-03-12T16:41:46.284691Z","shell.execute_reply":"2023-03-12T16:41:46.291085Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def load_dataset(filenames, labeled=True, ordered=False):\n    dataset = tf.data.TFRecordDataset(filenames)\n    dataset = dataset.map(read_tfrecord, num_parallel_calls=AUTOTUNE)\n    return dataset","metadata":{"execution":{"iopub.status.busy":"2023-03-12T16:41:46.295052Z","iopub.execute_input":"2023-03-12T16:41:46.295754Z","iopub.status.idle":"2023-03-12T16:41:46.302281Z","shell.execute_reply.started":"2023-03-12T16:41:46.295691Z","shell.execute_reply":"2023-03-12T16:41:46.301363Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"monet_ds = load_dataset(MONET_FILENAMES, labeled=True).batch(1)\nphoto_ds = load_dataset(PHOTO_FILENAMES, labeled=True).batch(1)","metadata":{"execution":{"iopub.status.busy":"2023-03-12T16:41:46.303643Z","iopub.execute_input":"2023-03-12T16:41:46.303993Z","iopub.status.idle":"2023-03-12T16:41:50.331074Z","shell.execute_reply.started":"2023-03-12T16:41:46.303967Z","shell.execute_reply":"2023-03-12T16:41:50.329935Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"example_monet = next(iter(monet_ds))\nexample_photo = next(iter(photo_ds))\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":"2023-03-12T16:41:50.332866Z","iopub.execute_input":"2023-03-12T16:41:50.333587Z","iopub.status.idle":"2023-03-12T16:41:51.182235Z","shell.execute_reply.started":"2023-03-12T16:41:50.333549Z","shell.execute_reply":"2023-03-12T16:41:51.181268Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Utilities","metadata":{}},{"cell_type":"code","source":"def create_gif(images_path, gif_path):\n    images = []\n    filenames = glob.glob(images_path)\n    print(filenames)\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        print(images)\n    imageio.mimsave(gif_path, images, fps=2) # Save gif","metadata":{"execution":{"iopub.status.busy":"2023-03-12T16:41:51.183349Z","iopub.execute_input":"2023-03-12T16:41:51.183711Z","iopub.status.idle":"2023-03-12T16:41:51.192225Z","shell.execute_reply.started":"2023-03-12T16:41:51.183674Z","shell.execute_reply":"2023-03-12T16:41:51.190896Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Build the generator¶\n\nWe'll be using a UNET architecture for our CycleGAN. To build our generator, let's first define our downsample and upsample methods.\n\nThe downsample, as the name suggests, reduces the 2D dimensions, the width and height, of the image by the stride. The stride is the length of the step the filter takes. Since the stride is 2, the filter is applied to every other pixel, hence reducing the weight and height by 2.\n\nWe'll be using an instance normalization instead of batch normalization. As the instance normalization is not standard in the TensorFlow API, we'll use the layer from TensorFlow Add-ons.\n","metadata":{}},{"cell_type":"code","source":"OUTPUT_CHANNELS = 3\n\ndef downsample(filters, size, apply_instancenorm=True):\n    initializer = tf.random_normal_initializer(0., 0.02)\n    gamma_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n\n    result = keras.Sequential()\n    result.add(layers.Conv2D(filters, size, strides=2, padding='same',\n                             kernel_initializer=initializer, use_bias=False))\n\n    if apply_instancenorm:\n        result.add(tfa.layers.InstanceNormalization(gamma_initializer=gamma_init))\n\n    result.add(layers.LeakyReLU())\n\n    return result\n\ndef downsample(filters, size, apply_instancenorm=True):\n    initializer = tf.random_normal_initializer(0., 0.02)\n\n    result = keras.Sequential()\n    result.add(layers.Conv2D(filters, size, strides=2, padding='same',\n                             kernel_initializer=initializer, use_bias=False))\n\n    if apply_instancenorm:\n        result.add(tfa.layers.InstanceNormalization())\n\n    result.add(layers.LeakyReLU())\n\n    return result","metadata":{"execution":{"iopub.status.busy":"2023-03-12T16:41:51.194538Z","iopub.execute_input":"2023-03-12T16:41:51.194981Z","iopub.status.idle":"2023-03-12T16:41:51.206460Z","shell.execute_reply.started":"2023-03-12T16:41:51.194927Z","shell.execute_reply":"2023-03-12T16:41:51.205326Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def upsample(filters, size, apply_dropout=False):\n    initializer = tf.random_normal_initializer(0., 0.03)\n    gamma_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n\n    result = keras.Sequential()\n    result.add(layers.Conv2DTranspose(filters, size, strides=2,\n                                      padding='same',\n                                      kernel_initializer=initializer,\n                                      use_bias=False))\n\n    result.add(tfa.layers.InstanceNormalization(gamma_initializer=gamma_init))\n\n    if apply_dropout:\n        result.add(layers.Dropout(0.5))\n\n    result.add(layers.ReLU())\n\n    return result\n\n# def upsample(filters, size, apply_dropout=False):\n#     initializer = tf.random_normal_initializer(0., 0.02)\n\n#     result = keras.Sequential()\n#     result.add(layers.Conv2DTranspose(filters, size, strides=2,\n#                                       padding='same',\n#                                       kernel_initializer=initializer,\n#                                       use_bias=False))\n#     result.add(tfa.layers.InstanceNormalization())\n\n#     if apply_dropout:\n#         result.add(layers.Dropout(0.5))\n\n#     result.add(layers.ReLU())\n\n#     return result","metadata":{"execution":{"iopub.status.busy":"2023-03-12T16:41:51.208244Z","iopub.execute_input":"2023-03-12T16:41:51.208930Z","iopub.status.idle":"2023-03-12T16:41:51.216881Z","shell.execute_reply.started":"2023-03-12T16:41:51.208897Z","shell.execute_reply":"2023-03-12T16:41:51.216295Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Let's build our generator!\n\nThe generator first downsamples the input image and then upsample while establishing long skip connections. Skip connections are a way to help bypass the vanishing gradient problem by concatenating the output of a layer to multiple layers instead of only one. Here we concatenate the output of the downsample layer to the upsample layer in a symmetrical fashion.\n","metadata":{}},{"cell_type":"code","source":"# def Generator():\n#     inputs = layers.Input(shape=[256,256,3])\n\n#     # bs = batch size\n#     down_stack = [\n#         downsample(64, 4, apply_instancenorm=False), # (bs, 128, 128, 64)\n#         downsample(128, 4), # (bs, 64, 64, 128)\n#         downsample(256, 4), # (bs, 32, 32, 256)\n#         downsample(512, 4), # (bs, 16, 16, 512)\n#         downsample(512, 4), # (bs, 8, 8, 512)\n#         downsample(512, 4), # (bs, 4, 4, 512)\n#         downsample(512, 4), # (bs, 2, 2, 512)\n#         downsample(512, 4), # (bs, 1, 1, 512)\n#     ]\n\n#     up_stack = [\n#         upsample(512, 4, apply_dropout=True), # (bs, 2, 2, 1024)\n#         upsample(512, 4, apply_dropout=True), # (bs, 4, 4, 1024)\n#         upsample(512, 4, apply_dropout=True), # (bs, 8, 8, 1024)\n#         upsample(512, 4), # (bs, 16, 16, 1024)\n#         upsample(256, 4), # (bs, 32, 32, 512)\n#         upsample(128, 4), # (bs, 64, 64, 256)\n#         upsample(64, 4), # (bs, 128, 128, 128)\n#     ]\n\n#     initializer = tf.random_normal_initializer(0., 0.02)\n#     last = layers.Conv2DTranspose(OUTPUT_CHANNELS, 4,\n#                                   strides=2,\n#                                   padding='same',\n#                                   kernel_initializer=initializer,\n#                                   activation='tanh') # (bs, 256, 256, 3)\n\n#     x = inputs\n\n#     # Downsampling through the model\n#     skips = []\n#     for down in down_stack:\n#         x = down(x)\n#         skips.append(x)\n\n#     skips = reversed(skips[:-1])\n\n#     # Upsampling and establishing the skip connections\n#     for up, skip in zip(up_stack, skips):\n#         x = up(x)\n#         x = layers.Concatenate()([x, skip])\n\n#     x = last(x)\n\n#     return keras.Model(inputs=inputs, outputs=x)\n\n\n\ndef make_generator():\n    inputs = layers.Input(shape=[256,256,3])\n\n    # bs = batch size\n    down_stack = [\n        downsample(64, 4, apply_instancenorm=False), # (bs, 128, 128, 64)\n        downsample(128, 4), # (bs, 64, 64, 128)\n        downsample(256, 4), # (bs, 32, 32, 256)\n        downsample(512, 4), # (bs, 16, 16, 512)\n        downsample(512, 4), # (bs, 8, 8, 512)\n        downsample(512, 4), # (bs, 4, 4, 512)\n        downsample(512, 4), # (bs, 2, 2, 512)\n        downsample(512, 4), # (bs, 1, 1, 512)\n    ]\n\n    up_stack = [\n        upsample(512, 4, apply_dropout=True), # (bs, 2, 2, 1024)\n        upsample(512, 4, apply_dropout=True), # (bs, 4, 4, 1024)\n        upsample(512, 4, apply_dropout=True), # (bs, 8, 8, 1024)\n        upsample(512, 4), # (bs, 16, 16, 1024)\n        upsample(256, 4), # (bs, 32, 32, 512)\n        upsample(128, 4), # (bs, 64, 64, 256)\n        upsample(64, 4), # (bs, 128, 128, 128)\n    ]\n    \n    # Change the kernel size of the last layer to 4, strides to 2, and set padding to \"same\"\n    initializer = tf.random_normal_initializer(0., 0.02)\n    last = layers.Conv2DTranspose(OUTPUT_CHANNELS, 4,\n                                  strides=2,\n                                  padding='same',\n                                  kernel_initializer=initializer,\n                                  activation='tanh') # (bs, 256, 256, 3)\n\n    x = inputs\n\n    # Downsampling through the model\n    skips = []\n    for down in down_stack:\n        x = down(x)\n        skips.append(x)\n\n    skips = reversed(skips[:-1])\n\n    # Upsampling and establishing the skip connections\n    for up, skip in zip(up_stack, skips):\n        x = up(x)\n        x = layers.Concatenate()([x, skip])\n\n    x = last(x)\n\n    return keras.Model(inputs=inputs, outputs=x)","metadata":{"execution":{"iopub.status.busy":"2023-03-12T16:41:51.221500Z","iopub.execute_input":"2023-03-12T16:41:51.223633Z","iopub.status.idle":"2023-03-12T16:41:51.235100Z","shell.execute_reply.started":"2023-03-12T16:41:51.223607Z","shell.execute_reply":"2023-03-12T16:41:51.234398Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Build the Discriminator\nThis Colab notebook code snippet defines the discriminator for PatchGAN, which is a type of convolutional neural network used for image-to-image translation tasks. PatchGAN aims to classify whether small patches of an image are real or fake by generating a classification map of the input image. The discriminator takes as input an image of shape (320, 320, 3), applies four 3x3 convolutional layers with stride 2 to downsample the image by a factor of 16, and then applies a final 3x3 convolutional layer with stride 1 to get a receptive field of 96x96. The output of the discriminator is a 20x20 classification map, where each element in the map represents the probability that the corresponding 96x96 patch in the input image is real.\n\nThe code uses the TensorFlow Keras API to define the discriminator model, with random normal weight initialization and a random normal gamma initialization for the batch normalization layers. The discriminator is trained with the Adam optimizer and binary cross-entropy loss to discriminate between real and generated images. Additionally, the code defines two generator models for transforming photos into Monet-esque paintings and Monet paintings into more photorealistic images, respectively. The discriminator models are used to distinguish between real and generated Monet paintings and photos. This PatchGAN discriminator is a critical component of the image-to-image translation pipeline and helps to improve the quality of the generated images.","metadata":{}},{"cell_type":"code","source":"def make_discriminator():\n    initializer = tf.random_normal_initializer(0., 0.02)\n    gamma_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n\n    inp = layers.Input(shape=[256, 256, 3], name='input_image')\n\n    # Apply 3x3 conv with stride 2 four times to downsample the image by a factor of 2^4 = 16\n    down1 = layers.Conv2D(64, kernel_size=3, strides=2, padding='same', kernel_initializer=initializer)(inp)  # (bs, 128, 128, 64)\n    down2 = layers.Conv2D(128, kernel_size=3, strides=2, padding='same', kernel_initializer=initializer)(down1)  # (bs, 64, 64, 128)\n    down3 = layers.Conv2D(256, kernel_size=3, strides=2, padding='same', kernel_initializer=initializer)(down2)  # (bs, 32, 32, 256)\n    down4 = layers.Conv2D(512, kernel_size=3, strides=2, padding='same', kernel_initializer=initializer)(down3)  # (bs, 16, 16, 512)\n\n    # Apply 3x3 conv with stride 1 to get a receptive field of 32x32\n    conv = layers.Conv2D(512, kernel_size=3, strides=1, padding='same', kernel_initializer=initializer)(down4)  # (bs, 16, 16, 512)\n    batchnorm1 = layers.BatchNormalization(gamma_initializer=gamma_init)(conv)\n    leaky_relu1 = layers.LeakyReLU(alpha=0.2)(batchnorm1)\n\n    # Apply 3x3 conv with stride 2 to get a receptive field of 64x64\n    conv2 = layers.Conv2D(512, kernel_size=3, strides=2, padding='same', kernel_initializer=initializer)(leaky_relu1)  # (bs, 8, 8, 512)\n    batchnorm2 = layers.BatchNormalization(gamma_initializer=gamma_init)(conv2)\n    leaky_relu2 = layers.LeakyReLU(alpha=0.2)(batchnorm2)\n\n    # Apply 3x3 conv with stride 1 to get a receptive field of 128x128\n    conv3 = layers.Conv2D(512, kernel_size=3, strides=1, padding='same', kernel_initializer=initializer)(leaky_relu2)  # (bs, 8, 8, 512)\n    batchnorm3 = layers.BatchNormalization(gamma_initializer=gamma_init)(conv3)\n    leaky_relu3 = layers.LeakyReLU(alpha=0.2)(batchnorm3)\n\n    # Apply 3x3 conv with stride 2 to get a receptive field of 256x256\n    patch = layers.Conv2D(1, kernel_size=3, strides=2, padding='same', kernel_initializer=initializer)(leaky_relu3)  # (bs, 4, 4, 1)\n\n    return tf.keras.Model(inputs=inp, outputs=patch)\n\nwith strategy.scope():\n    monet_generator = make_generator() # transforms photos to Monet-esque paintings\n    photo_generator = make_generator() # transforms Monet paintings to be more like photos\n\n    monet_discriminator = make_discriminator() # differentiates real Monet paintings and generated Monet paintings\n    photo_discriminator = make_discriminator() # differentiates real photos and generated photos","metadata":{"execution":{"iopub.status.busy":"2023-03-12T16:41:51.236398Z","iopub.execute_input":"2023-03-12T16:41:51.237018Z","iopub.status.idle":"2023-03-12T16:41:53.589043Z","shell.execute_reply.started":"2023-03-12T16:41:51.236984Z","shell.execute_reply":"2023-03-12T16:41:53.587998Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Since our generators are not trained yet, the generated Monet-esque photo does not show what is expected at this point.","metadata":{}},{"cell_type":"code","source":"to_monet = monet_generator(example_photo)\n\nplt.subplot(1, 2, 1)\nplt.title(\"Original Photo\")\nplt.imshow(example_photo[0] * 0.5 + 0.5)\n\nplt.subplot(1, 2, 2)\nplt.title(\"Monet-esque Photo\")\nplt.imshow(to_monet[0] * 0.5 + 0.5)\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2023-03-12T16:41:53.590551Z","iopub.execute_input":"2023-03-12T16:41:53.590897Z","iopub.status.idle":"2023-03-12T16:42:02.524706Z","shell.execute_reply.started":"2023-03-12T16:41:53.590862Z","shell.execute_reply":"2023-03-12T16:42:02.523617Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Build the CycleGAN model\n\nWe will subclass a `tf.keras.Model` so that we can run `fit()` later to train our model. During the training step, the model transforms a photo to a Monet painting and then back to a photo. The difference between the original photo and the twice-transformed photo is the cycle-consistency loss. We want the original photo and the twice-transformed photo to be similar to one another.\n\nThe losses are defined in the next section.","metadata":{}},{"cell_type":"code","source":"class CycleGan(keras.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        monet_gen_optimizer,\n        photo_gen_optimizer,\n        monet_disc_optimizer,\n        photo_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 = monet_gen_optimizer\n        self.p_gen_optimizer = photo_gen_optimizer\n        self.m_disc_optimizer = monet_disc_optimizer\n        self.p_disc_optimizer = photo_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 {\n            \"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":"2023-03-12T16:42:02.526721Z","iopub.execute_input":"2023-03-12T16:42:02.527585Z","iopub.status.idle":"2023-03-12T16:42:02.546511Z","shell.execute_reply.started":"2023-03-12T16:42:02.527540Z","shell.execute_reply":"2023-03-12T16:42:02.545101Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Callbacks - GANMonitor","metadata":{}},{"cell_type":"code","source":"# GANMonitor\nfrom tensorflow.keras.callbacks import Callback\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.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.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":{"execution":{"iopub.status.busy":"2023-03-12T16:42:02.548570Z","iopub.execute_input":"2023-03-12T16:42:02.549069Z","iopub.status.idle":"2023-03-12T16:42:02.565692Z","shell.execute_reply.started":"2023-03-12T16:42:02.549025Z","shell.execute_reply":"2023-03-12T16:42:02.564459Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Define loss functions\nThe discriminator loss function below compares real images to a matrix of 1s and fake images to a matrix of 0s. The perfect discriminator will output all 1s for real images and all 0s for fake images. The discriminator loss outputs the average of the real and generated los","metadata":{}},{"cell_type":"code","source":"with strategy.scope():\n    def discriminator_loss(real, generated):\n        real_loss = tf.keras.losses.BinaryCrossentropy(from_logits=True, reduction=tf.keras.losses.Reduction.NONE)(tf.ones_like(real), real)\n\n        generated_loss = tf.keras.losses.BinaryCrossentropy(from_logits=True, reduction=tf.keras.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","metadata":{"execution":{"iopub.status.busy":"2023-03-12T16:42:02.569240Z","iopub.execute_input":"2023-03-12T16:42:02.570447Z","iopub.status.idle":"2023-03-12T16:42:02.582407Z","shell.execute_reply.started":"2023-03-12T16:42:02.570405Z","shell.execute_reply":"2023-03-12T16:42:02.581155Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The generator wants to fool the discriminator into thinking the generated image is real. The perfect generator will have the discriminator output only 1s. Thus, it compares the generated image to a matrix of 1s to find the loss.","metadata":{}},{"cell_type":"code","source":"with strategy.scope():\n    def generator_loss(generated):\n        return tf.keras.losses.BinaryCrossentropy(from_logits=True, reduction=tf.keras.losses.Reduction.NONE)(tf.ones_like(generated), generated)","metadata":{"execution":{"iopub.status.busy":"2023-03-12T16:42:02.583873Z","iopub.execute_input":"2023-03-12T16:42:02.584692Z","iopub.status.idle":"2023-03-12T16:42:02.593752Z","shell.execute_reply.started":"2023-03-12T16:42:02.584650Z","shell.execute_reply":"2023-03-12T16:42:02.592337Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We want our original photo and the twice transformed photo to be similar to one another. Thus, we can calculate the cycle consistency loss be finding the average of their difference.\n\nThe identity loss compares the image with its generator (i.e. photo with photo generator). If given a photo as input, we want it to generate the same image as the image was originally a photo. The identity loss compares the input with the output of the generator.\n","metadata":{}},{"cell_type":"code","source":"with strategy.scope():\n    def calc_cycle_loss(real_image, cycled_image, LAMBDA):\n        loss = tf.reduce_mean(tf.abs(real_image - cycled_image))\n\n        return LAMBDA * loss\n        \nwith 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":"2023-03-12T16:42:02.595536Z","iopub.execute_input":"2023-03-12T16:42:02.596538Z","iopub.status.idle":"2023-03-12T16:42:02.608208Z","shell.execute_reply.started":"2023-03-12T16:42:02.596464Z","shell.execute_reply":"2023-03-12T16:42:02.607110Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Train the CycleGAN\n\nLet's compile our model. Since we used `tf.keras.Model` to build our CycleGAN, we can just ude the `fit` function to train our model.","metadata":{}},{"cell_type":"code","source":"with strategy.scope():\n    monet_generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n    photo_generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n\n    monet_discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n    photo_discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n\nwith strategy.scope():\n    cycle_gan_model = CycleGan(\n        monet_generator, photo_generator, monet_discriminator, photo_discriminator\n    )\n\n    cycle_gan_model.compile(\n        monet_gen_optimizer = monet_generator_optimizer,\n        photo_gen_optimizer = photo_generator_optimizer,\n        monet_disc_optimizer = monet_discriminator_optimizer,\n        photo_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\n    )\n\nhistory = cycle_gan_model.fit(\n    tf.data.Dataset.zip((monet_ds, photo_ds)),\n    callbacks=[GANMonitor()],\n    epochs=100\n)","metadata":{"execution":{"iopub.status.busy":"2023-03-12T16:42:02.610833Z","iopub.execute_input":"2023-03-12T16:42:02.611162Z","iopub.status.idle":"2023-03-12T18:09:12.236042Z","shell.execute_reply.started":"2023-03-12T16:42:02.611134Z","shell.execute_reply":"2023-03-12T18:09:12.234188Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Visualize our Monet-esque photos","metadata":{}},{"cell_type":"code","source":"_, ax = plt.subplots(5, 2, figsize=(12, 12))\nfor i, img in enumerate(photo_ds.skip(10).take(5)):\n    prediction = monet_generator(img, training=False)[0].numpy()\n    prediction = (prediction * 127.5 + 127.5).astype(np.uint8)\n    img = (img[0] * 127.5 + 127.5).numpy().astype(np.uint8)\n\n    ax[i, 0].imshow(img)\n    ax[i, 1].imshow(prediction)\n    ax[i, 0].set_title(\"Input Photo\")\n    ax[i, 1].set_title(\"Monet-esque\")\n    ax[i, 0].axis(\"off\")\n    ax[i, 1].axis(\"off\")\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2023-03-12T18:09:18.738736Z","iopub.execute_input":"2023-03-12T18:09:18.739090Z","iopub.status.idle":"2023-03-12T18:09:19.737910Z","shell.execute_reply.started":"2023-03-12T18:09:18.739058Z","shell.execute_reply":"2023-03-12T18:09:19.737086Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Submit","metadata":{}},{"cell_type":"code","source":"import PIL\n! mkdir ../images\n\ni = 1\nfor img in photo_ds:\n    prediction = monet_generator(img, training=False)[0].numpy()\n    prediction = (prediction * 127.5 + 127.5).astype(np.uint8)\n    im = PIL.Image.fromarray(prediction)\n    im.save(\"../images/\" + str(i) + \".jpg\")\n    i += 1\n    \nimport shutil\nshutil.make_archive(\"/kaggle/working/images\", 'zip', \"/kaggle/images\")","metadata":{"execution":{"iopub.status.busy":"2023-03-12T18:09:33.515144Z","iopub.execute_input":"2023-03-12T18:09:33.515558Z","iopub.status.idle":"2023-03-12T18:15:01.251588Z","shell.execute_reply.started":"2023-03-12T18:09:33.515522Z","shell.execute_reply":"2023-03-12T18:15:01.250284Z"},"trusted":true},"execution_count":null,"outputs":[]}]}