{"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\n\nUSING_KAGGLE = True\nAUTOTUNE = tf.data.experimental.AUTOTUNE\nBATCH_SIZE = 1\n\ntry:\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\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()\n","metadata":{"execution":{"iopub.status.busy":"2023-02-20T03:10:21.172039Z","iopub.execute_input":"2023-02-20T03:10:21.172551Z","iopub.status.idle":"2023-02-20T03:10:26.238911Z","shell.execute_reply.started":"2023-02-20T03:10:21.172520Z","shell.execute_reply":"2023-02-20T03:10:26.238241Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from typing import List\n\nimport tensorflow as tf\n\n\ndef decode_image(image: tf.io.FixedLenFeature, image_size: List[int] = [256, 256]) -> tf.Tensor:\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\n\ndef read_tfrecord(example: tf.Tensor) -> tf.Tensor:\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\n\n\ndef read_tfrecords(filenames: 'list[str]') -> tf.data.TFRecordDataset:\n    dataset = tf.data.TFRecordDataset(filenames)\n    dataset = dataset.map(read_tfrecord, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n    return dataset\n","metadata":{"execution":{"iopub.status.busy":"2023-02-20T03:10:26.239868Z","iopub.execute_input":"2023-02-20T03:10:26.240391Z","iopub.status.idle":"2023-02-20T03:10:26.248852Z","shell.execute_reply.started":"2023-02-20T03:10:26.240360Z","shell.execute_reply":"2023-02-20T03:10:26.248125Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from typing import Callable\n\nimport tensorflow as tf\nimport numpy as np\nimport re\n\n\ndef dataset_path(path: str = '../Sample_Data/Raw') -> str:\n    \"\"\"This function loads the Monet and photo images in `TFRecordDataset` format.\"\"\"\n    if USING_KAGGLE:\n        from kaggle_datasets import KaggleDatasets\n        return KaggleDatasets().get_gcs_path()\n    else:\n        return path\n\n\ndef monet_filenames():\n    \"\"\"This sub-function returns the filenames for the monet painting\"\"\"\n    return tf.io.gfile.glob(f'{dataset_path()}/monet_tfrec/*.tfrec')\n\n\ndef photo_filenames():\n    \"\"\"This sub-function returns the filenames for the photos\"\"\"\n    return tf.io.gfile.glob(f'{dataset_path()}/photo_tfrec/*.tfrec')\n\n\ndef monet_dataset():\n    \"\"\"This sub-function loads the Monet paintings in `TFRecordDataset` format\"\"\"\n    return read_tfrecords(monet_filenames())\n\n\ndef photo_dataset():\n    \"\"\"This sub-function loads the photos in `TFRecordDataset` format\"\"\"\n    return read_tfrecords(photo_filenames())\n\n\ndef count_data_items(filenames) -> np.array:\n    n: list\n\n    n = [int(re.compile(r\"-([0-9]*)\\.\").search(filename).group(1)) for filename in filenames]\n\n    return np.sum(n)\n\n\ndef count_monet_samples():\n    return count_data_items(monet_filenames())\n\n\ndef count_photo_samples():\n    return count_data_items(photo_filenames())\n\n\ndef get_gan_dataset_basic(batch_size: int = 1) -> tf.data.Dataset:\n    monets = monet_dataset().batch(batch_size, drop_remainder=True)\n    photos = photo_dataset().batch(batch_size, drop_remainder=True)\n    return tf.data.Dataset.zip((monets, photos))\n","metadata":{"execution":{"iopub.status.busy":"2023-02-20T03:10:26.249859Z","iopub.execute_input":"2023-02-20T03:10:26.250439Z","iopub.status.idle":"2023-02-20T03:10:26.264083Z","shell.execute_reply.started":"2023-02-20T03:10:26.250407Z","shell.execute_reply":"2023-02-20T03:10:26.263222Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import tensorflow as tf\nfrom tensorflow_addons.layers import InstanceNormalization\n\n\ndef downsample(filters: int,\n               size: int,\n               apply_instancenorm: bool = True) -> tf.keras.Sequential:\n    initializer = tf.random_normal_initializer(0., 0.02)\n    gamma_init = tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n\n    result = tf.keras.Sequential()\n    result.add(tf.keras.layers.Conv2D(filters,\n                                      size,\n                                      strides=2,\n                                      padding='same',\n                                      kernel_initializer=initializer,\n                                      use_bias=False))\n    if apply_instancenorm:\n        result.add(InstanceNormalization(gamma_initializer=gamma_init))\n    result.add(tf.keras.layers.ReLU())\n\n    return result\n\n\ndef upsample(filters: int,\n             size: int,\n             apply_dropout: bool = False) -> tf.keras.Sequential:\n    initializer = tf.random_normal_initializer(0., 0.02)\n    gamma_init = tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n\n    result = tf.keras.Sequential()\n    result.add(tf.keras.layers.Conv2DTranspose(filters,\n                                               size,\n                                               strides=2,\n                                               padding='same',\n                                               kernel_initializer=initializer,\n                                               use_bias=False))\n    result.add(InstanceNormalization(gamma_initializer=gamma_init))\n    if apply_dropout:\n        result.add(tf.keras.layers.Dropout(0.5))\n    result.add(tf.keras.layers.ReLU())\n\n    return result\n","metadata":{"execution":{"iopub.status.busy":"2023-02-20T03:10:26.265905Z","iopub.execute_input":"2023-02-20T03:10:26.266270Z","iopub.status.idle":"2023-02-20T03:10:26.280509Z","shell.execute_reply.started":"2023-02-20T03:10:26.266237Z","shell.execute_reply":"2023-02-20T03:10:26.279271Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import tensorflow as tf\nfrom typing import List\n\n\ndef generator(output_channels: int = 3) -> tf.keras.Model:\n    \"\"\"\n    The Generator applies an encoder-decoder architecture that downsamples the image and then decodes the image\n    through various transpose convolutions.\" It is inspired by a \"U-Net\" architecture for image generation, which downsamples\n    an image and then applies \n    \"\"\"\n    down_stack = [\n        downsample(64, 4, apply_instancenorm=False),\n        downsample(128, 4),\n        downsample(256, 4),\n        downsample(512, 4),\n        downsample(512, 4),\n        downsample(512, 4),\n        downsample(512, 4),\n        downsample(512, 4),\n    ]\n\n    up_stack = [\n        upsample(512, 4, apply_dropout=True),\n        upsample(512, 4, apply_dropout=True),\n        upsample(512, 4, apply_dropout=True),\n        upsample(512, 4),\n        upsample(256, 4),\n        upsample(128, 4),\n        upsample(64, 4),\n    ]\n\n    initializer = tf.random_normal_initializer(0., 0.02)\n\n    inputs = tf.keras.layers.Input(shape=[256, 256, 3])\n    x = inputs\n\n    skips: List[tf.keras.Sequential] = []\n    for down in down_stack:\n        x = down(x)\n        skips.append(x)\n\n    skips = list(reversed(skips[:-1]))\n\n    for up, skip in zip(up_stack, skips):\n        x = up(x)\n        x = tf.keras.layers.Concatenate()([x, skip])\n\n    x = tf.keras.layers.Conv2DTranspose(output_channels,\n                                        4,\n                                        strides=2,\n                                        padding='same',\n                                        kernel_initializer=initializer,\n                                        activation='tanh')(x)\n\n    return tf.keras.Model(inputs=inputs, outputs=x)\n","metadata":{"execution":{"iopub.status.busy":"2023-02-20T03:10:26.282021Z","iopub.execute_input":"2023-02-20T03:10:26.282304Z","iopub.status.idle":"2023-02-20T03:10:26.296792Z","shell.execute_reply.started":"2023-02-20T03:10:26.282274Z","shell.execute_reply":"2023-02-20T03:10:26.295809Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import tensorflow as tf\nfrom tensorflow_addons.layers import InstanceNormalization\n\n\ndef discriminator() -> tf.keras.Model:\n    initializer = tf.random_normal_initializer(0., 0.02)\n    gamma_init = tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n\n    inputs = tf.keras.layers.Input(shape=[256, 256, 3], name='input_image')\n\n    x = inputs\n    down1 = downsample(64, 4, False)(x)\n    down2 = downsample(128, 4)(down1)\n    down3 = downsample(256, 4)(down2)\n    zero_pad1 = tf.keras.layers.ZeroPadding2D()(down3)\n    conv = tf.keras.layers.Conv2D(512, 4, strides=1, kernel_initializer=initializer, use_bias=False)(zero_pad1)\n    norm1 = InstanceNormalization(gamma_initializer=gamma_init)(conv)\n    leaky_relu = tf.keras.layers.LeakyReLU()(norm1)\n    zero_pad2 = tf.keras.layers.ZeroPadding2D()(leaky_relu)\n    outputs = tf.keras.layers.Conv2D(1, 4, strides=1, kernel_initializer=initializer)(zero_pad2)\n\n    return tf.keras.Model(inputs=inputs, outputs=outputs)\n","metadata":{"execution":{"iopub.status.busy":"2023-02-20T03:10:26.298284Z","iopub.execute_input":"2023-02-20T03:10:26.298773Z","iopub.status.idle":"2023-02-20T03:10:26.311204Z","shell.execute_reply.started":"2023-02-20T03:10:26.298732Z","shell.execute_reply":"2023-02-20T03:10:26.310211Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from typing import Callable, Dict, List\n\nimport tensorflow as tf\n\n\nwith strategy.scope():\n    def diff_augment(x: tf.Tensor,\n                     policy: str = '',\n                     channels_first: bool = False) -> tf.Tensor:\n        if policy:\n            if channels_first:\n                x = tf.transpose(x, [0, 2, 3, 1])\n            for p in policy.split(','):\n                for fn in AUGMENT_FNS[p]:\n                    x = fn(x)\n            if channels_first:\n                x = tf.transpose(x, [0, 3, 1, 2])\n        return x\n\n\n    def rand_brightness(x: tf.Tensor) -> tf.Tensor:\n        magnitude: tf.Tensor\n\n        magnitude = tf.random.uniform([tf.shape(x)[0], 1, 1, 1]) - 0.5\n        x = x + magnitude\n\n        return x\n\n\n    def rand_saturation(x: tf.Tensor) -> tf.Tensor:\n        magnitude: tf.Tensor\n        x_mean: tf.Tensor\n\n        magnitude = tf.random.uniform([tf.shape(x)[0], 1, 1, 1]) * 2\n        x_mean = tf.reduce_mean(x, axis=3, keepdims=True) * 0.3333333333333333333\n        x = (x - x_mean) * magnitude + x_mean\n\n        return x\n\n\n    def rand_contrast(x: tf.Tensor) -> tf.Tensor:\n        magnitude: tf.Tensor\n        x_mean: tf.Tensor\n\n        magnitude = tf.random.uniform([tf.shape(x)[0], 1, 1, 1]) + 0.5\n        x_mean = tf.reduce_mean(x, axis=[1, 2, 3], keepdims=True) * 5.086e-6\n        x = (x - x_mean) * magnitude + x_mean\n\n        return x\n\n\n    def rand_translation(x: tf.Tensor,\n                         ratio: float = 0.125) -> tf.Tensor:\n        batch_size: tf.Tensor\n        image_size: tf.Tensor\n        shift: tf.Tensor\n        translation_x: tf.Tensor\n        translation_y: tf.Tensor\n        grid_x: tf.Tensor\n        grid_y: tf.Tensor\n\n        batch_size = tf.shape(x)[0]\n        image_size = tf.shape(x)[1:3]\n        shift = tf.cast(tf.cast(image_size, tf.float32) * ratio + 0.5, tf.int32)\n        translation_x = tf.random.uniform([batch_size, 1], -shift[0], shift[0] + 1, dtype=tf.int32)\n        translation_y = tf.random.uniform([batch_size, 1], -shift[1], shift[1] + 1, dtype=tf.int32)\n        grid_x = tf.clip_by_value(tf.expand_dims(tf.range(image_size[0], dtype=tf.int32), 0) + translation_x + 1,\n                                  0,\n                                  image_size[0] + 1)\n        grid_y = tf.clip_by_value(tf.expand_dims(tf.range(image_size[1], dtype=tf.int32), 0) + translation_y + 1,\n                                  0,\n                                  image_size[1] + 1)\n        x = tf.gather_nd(tf.pad(x, [[0, 0], [1, 1], [0, 0], [0, 0]]), tf.expand_dims(grid_x, -1), batch_dims=1)\n        x = tf.transpose(tf.gather_nd(tf.pad(tf.transpose(x, [0, 2, 1, 3]), [[0, 0], [1, 1], [0, 0], [0, 0]]),\n                                      tf.expand_dims(grid_y, -1),\n                                      batch_dims=1),\n                         [0, 2, 1, 3])\n\n        return x\n\n\n    def rand_cutout(x: tf.Tensor,\n                    ratio: float = 0.5) -> tf.Tensor:\n        batch_size: tf.Tensor\n        image_size: tf.Tensor\n        cutout_size: tf.Tensor\n        offset_x: tf.Tensor\n        offset_y: tf.Tensor\n        grid_batch: tf.Tensor\n        grid_x: tf.Tensor\n        grid_y: tf.Tensor\n        cutout_grid: tf.Tensor\n        mask: tf.Tensor\n\n        batch_size = tf.shape(x)[0]\n        image_size = tf.shape(x)[1:3]\n        cutout_size = tf.cast(tf.cast(image_size, tf.float32) * ratio + 0.5, tf.int32)\n        offset_x = tf.random.uniform([tf.shape(x)[0], 1, 1],\n                                     maxval=image_size[0] + (1 - cutout_size[0] % 2),\n                                     dtype=tf.int32)\n        offset_y = tf.random.uniform([tf.shape(x)[0], 1, 1],\n                                     maxval=image_size[1] + (1 - cutout_size[1] % 2),\n                                     dtype=tf.int32)\n        grid_batch, grid_x, grid_y = tf.meshgrid(tf.range(batch_size, dtype=tf.int32),\n                                                 tf.range(cutout_size[0], dtype=tf.int32),\n                                                 tf.range(cutout_size[1], dtype=tf.int32),\n                                                 indexing='ij')\n        cutout_grid = tf.stack([grid_batch,\n                                grid_x + offset_x - cutout_size[0] // 2,\n                                grid_y + offset_y - cutout_size[1] // 2],\n                               axis=-1)\n        mask_shape = tf.stack([batch_size, image_size[0], image_size[1]])\n        cutout_grid = tf.maximum(cutout_grid, 0)\n        cutout_grid = tf.minimum(cutout_grid, tf.reshape(mask_shape - 1, [1, 1, 1, 3]))\n        mask = tf.maximum(1 - tf.scatter_nd(cutout_grid,\n                                            tf.ones([batch_size, cutout_size[0], cutout_size[1]], dtype=tf.float32),\n                                            mask_shape),\n                          0)\n        x = x * tf.expand_dims(mask, axis=3)\n\n        return x\n\n\n    AUGMENT_FNS: Dict[str, List[Callable[[tf.Tensor], tf.Tensor]]] = {\n        'color': [rand_brightness, rand_saturation, rand_contrast],\n        'translation': [rand_translation],\n        'cutout': [rand_cutout],\n    }\n\n\n    def aug_fn(image: tf.Tensor) -> tf.Tensor:\n        return diff_augment(image, \"color,translation,cutout\")\n\n\ndef data_augment_flip(image: tf.Tensor) -> tf.Tensor:\n    image = tf.image.random_flip_left_right(image)\n    return image\n","metadata":{"execution":{"iopub.status.busy":"2023-02-20T03:10:26.312761Z","iopub.execute_input":"2023-02-20T03:10:26.313483Z","iopub.status.idle":"2023-02-20T03:10:26.345740Z","shell.execute_reply.started":"2023-02-20T03:10:26.313445Z","shell.execute_reply":"2023-02-20T03:10:26.344647Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import tensorflow as tf\n\n\nwith strategy.scope():\n    def discriminator_loss(real: tf.keras.Model, generated: tf.keras.Model) -> tf.Tensor:\n        real_loss = tf.keras.losses.BinaryCrossentropy(from_logits=True,\n                                                       reduction=tf.keras.losses.Reduction.NONE)(tf.ones_like(real),\n                                                                                                 real)\n        generated_loss = tf.keras.losses.BinaryCrossentropy(from_logits=True,\n                                                            reduction=tf.keras.losses.Reduction.NONE)(\n            tf.zeros_like(generated), generated)\n        total_discriminator_loss = real_loss + generated_loss\n        return total_discriminator_loss * 0.5\n\n\n    def generator_loss(generated: tf.keras.Model) -> tf.Tensor:\n        return tf.keras.losses.BinaryCrossentropy(from_logits=True,\n                                                  reduction=tf.keras.losses.Reduction.NONE)(tf.ones_like(generated),\n                                                                                            generated)\n\n\n    def calc_cycle_loss(real_image: tf.Tensor, cycled_image: tf.Tensor, alpha: float) -> tf.Tensor:\n        loss1: tf.Tensor = tf.reduce_mean(tf.abs(real_image - cycled_image))\n        return loss1 * alpha\n\n\n    def identity_loss(real_image: tf.Tensor, same_image: tf.Tensor, alpha: float) -> tf.Tensor:\n        loss: tf.Tensor = tf.reduce_mean(tf.abs(real_image - same_image))\n        return alpha * 0.5 * loss\n","metadata":{"execution":{"iopub.status.busy":"2023-02-20T03:10:26.346930Z","iopub.execute_input":"2023-02-20T03:10:26.347222Z","iopub.status.idle":"2023-02-20T03:10:26.362426Z","shell.execute_reply.started":"2023-02-20T03:10:26.347195Z","shell.execute_reply":"2023-02-20T03:10:26.361385Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from abc import ABC\nfrom typing import Callable, Tuple\n\nimport tensorflow as tf\n\n\nclass CycleGan(tf.keras.Model, ABC):\n    \"\"\"\n    This class is the CycleGAN model, which initializes and compiles the generators and discriminators\n    \"\"\"\n\n    def __init__(self,\n                 monet_generator: tf.keras.Model,\n                 photo_generator: tf.keras.Model,\n                 monet_discriminator: tf.keras.Model,\n                 photo_discriminator: tf.keras.Model,\n                 lambda_cycle: int = 10):\n        \"\"\"\n        This is the initialization function for the generators, discriminators, and the lambda cycle\n        \"\"\"\n        super(CycleGan, self).__init__()\n        self.monet_generator = monet_generator\n        self.photo_generator = photo_generator\n        self.monet_discriminator = monet_discriminator\n        self.photo_discriminator = photo_discriminator\n        self.lambda_cycle = lambda_cycle\n        self.monet_generator_optimizer: tf.keras.optimizers.Optimizer = None\n        self.photo_generator_optimizer: tf.keras.optimizers.Optimizer = None\n        self.monet_discriminator_optimizer: tf.keras.optimizers.Optimizer = None\n        self.photo_discriminator_optimizer: tf.keras.optimizers.Optimizer = None\n        self.generator_loss_fn: Callable[[tf.keras.Model], tf.Tensor] = None\n        self.discriminator_loss_fn: Callable[[tf.keras.Model, tf.keras.Model], tf.Tensor] = None\n        self.cycle_loss_fn: Callable[[tf.Tensor, tf.Tensor, float], tf.Tensor] = None\n        self.identity_loss_fn: Callable[[tf.Tensor, tf.Tensor, float], tf.Tensor] = None\n        self.aug_fn = None\n\n    def compile(self,\n                monet_generator_optimizer: tf.keras.optimizers.Optimizer,\n                photo_generator_optimizer: tf.keras.optimizers.Optimizer,\n                monet_discriminator_optimizer: tf.keras.optimizers.Optimizer,\n                photo_discriminator_optimizer: tf.keras.optimizers.Optimizer,\n                generator_loss_fn: Callable[[tf.keras.Model], tf.Tensor],\n                discriminator_loss_fn: Callable[[tf.keras.Model, tf.keras.Model], tf.Tensor],\n                cycle_loss_fn: Callable[[tf.Tensor, tf.Tensor, float], tf.Tensor],\n                identity_loss_fn: Callable[[tf.Tensor, tf.Tensor, float], tf.Tensor],\n                aug_fn):\n        \"\"\"\n        This function sets the optimizers and the loss functions\n        \"\"\"\n        super(CycleGan, self).compile()\n        self.monet_generator_optimizer = monet_generator_optimizer\n        self.photo_generator_optimizer = photo_generator_optimizer\n        self.monet_discriminator_optimizer = monet_discriminator_optimizer\n        self.photo_discriminator_optimizer = photo_discriminator_optimizer\n        self.generator_loss_fn = generator_loss_fn\n        self.discriminator_loss_fn = discriminator_loss_fn\n        self.cycle_loss_fn = cycle_loss_fn\n        self.identity_loss_fn = identity_loss_fn\n        self.aug_fn = aug_fn\n\n    @tf.function\n    def train_step(self, batch_data: Tuple[tf.Tensor, tf.Tensor]):\n        real_monet, real_photo = batch_data\n        batch_size = tf.shape(real_monet)[0]\n\n        with tf.GradientTape(persistent=True) as tape:\n            fake_monet = self.monet_generator(real_photo, training=True)\n            cycled_photo = self.photo_generator(fake_monet, training=True)\n\n            fake_photo = self.photo_generator(real_monet, training=True)\n            cycled_monet = self.monet_generator(fake_photo, training=True)\n\n            same_monet = self.monet_generator(real_monet, training=True)\n            same_photo = self.photo_generator(real_photo, training=True)\n\n            both_monet = tf.concat([real_monet, fake_monet], axis=0)\n\n            aug_monet = aug_fn(both_monet)\n\n            aug_real_monet = aug_monet[:batch_size]\n            aug_fake_monet = aug_monet[batch_size:]\n\n            discriminator_real_monet = self.monet_discriminator(aug_real_monet, training=True)\n            discriminator_real_photo = self.photo_discriminator(real_photo, training=True)\n\n            discriminator_fake_monet = self.monet_discriminator(aug_fake_monet, training=True)\n            discriminator_fake_photo = self.photo_discriminator(fake_photo, training=True)\n\n            monet_generator_loss = self.generator_loss_fn(discriminator_fake_monet)\n            photo_generator_loss = self.generator_loss_fn(discriminator_fake_photo)\n\n            total_cycle_loss = (self.cycle_loss_fn(real_monet,\n                                                   cycled_monet,\n                                                   self.lambda_cycle)\n                                + self.cycle_loss_fn(real_photo,\n                                                     cycled_photo,\n                                                     self.lambda_cycle))\n\n            total_monet_generator_loss = (monet_generator_loss\n                                          + total_cycle_loss\n                                          + self.identity_loss_fn(real_monet,\n                                                                  same_monet,\n                                                                  self.lambda_cycle))\n            total_photo_generator_loss = (photo_generator_loss\n                                          + total_cycle_loss\n                                          + self.identity_loss_fn(real_photo,\n                                                                  same_photo,\n                                                                  self.lambda_cycle))\n\n            monet_discriminator_loss = self.discriminator_loss_fn(discriminator_real_monet, discriminator_fake_monet)\n            photo_discriminator_loss = self.discriminator_loss_fn(discriminator_real_photo, discriminator_fake_photo)\n\n        monet_generator_gradients = tape.gradient(total_monet_generator_loss, self.monet_generator.trainable_variables)\n        photo_generator_gradients = tape.gradient(total_photo_generator_loss, self.photo_generator.trainable_variables)\n\n        monet_discriminator_gradients = tape.gradient(monet_discriminator_loss,\n                                                      self.monet_discriminator.trainable_variables)\n        photo_discriminator_gradients = tape.gradient(photo_discriminator_loss,\n                                                      self.photo_discriminator.trainable_variables)\n\n        self.monet_generator_optimizer.apply_gradients(zip(monet_generator_gradients,\n                                                           self.monet_generator.trainable_variables))\n        self.photo_generator_optimizer.apply_gradients(zip(photo_generator_gradients,\n                                                           self.photo_generator.trainable_variables))\n\n        self.monet_discriminator_optimizer.apply_gradients(zip(monet_discriminator_gradients,\n                                                               self.monet_discriminator.trainable_variables))\n        self.photo_discriminator_optimizer.apply_gradients(zip(photo_discriminator_gradients,\n                                                               self.photo_discriminator.trainable_variables))\n\n        return {\n            'monet_generator_loss': total_monet_generator_loss,\n            'photo_generator_loss': total_photo_generator_loss,\n            'monet_discriminator_loss': monet_discriminator_loss,\n            'photo_discriminator_loss': photo_discriminator_loss\n        }\n","metadata":{"execution":{"iopub.status.busy":"2023-02-20T03:10:26.363796Z","iopub.execute_input":"2023-02-20T03:10:26.364104Z","iopub.status.idle":"2023-02-20T03:10:26.391424Z","shell.execute_reply.started":"2023-02-20T03:10:26.364046Z","shell.execute_reply":"2023-02-20T03:10:26.390530Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import tensorflow as tf\n\n\nmonet_generator = generator()\nphoto_generator = generator()\nmonet_discriminator = discriminator()\nphoto_discriminator = discriminator()\n\ncycle_gan_model = CycleGan(monet_generator=monet_generator,\n                           photo_generator=photo_generator,\n                           monet_discriminator=monet_discriminator,\n                           photo_discriminator=photo_discriminator)\n\nwith strategy.scope():\n    def cycle_gan_compile(monet_generator_optimizer: tf.keras.optimizers.Optimizer,\n                          photo_generator_optimizer: tf.keras.optimizers.Optimizer,\n                          monet_discriminator_optimizer: tf.keras.optimizers.Optimizer,\n                          photo_discriminator_optimizer: tf.keras.optimizers.Optimizer) -> None:\n        cycle_gan_model.compile(\n            monet_generator_optimizer=monet_generator_optimizer,\n            photo_generator_optimizer=photo_generator_optimizer,\n            monet_discriminator_optimizer=monet_discriminator_optimizer,\n            photo_discriminator_optimizer=photo_discriminator_optimizer,\n            generator_loss_fn=generator_loss,\n            discriminator_loss_fn=discriminator_loss,\n            cycle_loss_fn=calc_cycle_loss,\n            identity_loss_fn=identity_loss,\n            aug_fn=aug_fn\n        )\n\n\n    def cycle_gan_compile_with_loss_rate(loss_rate: float) -> None:\n        monet_generator_optimizer = tf.keras.optimizers.Adam(loss_rate, beta_1=0.5)\n        photo_generator_optimizer = tf.keras.optimizers.Adam(loss_rate, beta_1=0.5)\n        monet_discriminator_optimizer = tf.keras.optimizers.Adam(loss_rate, beta_1=0.5)\n        photo_discriminator_optimizer = tf.keras.optimizers.Adam(loss_rate, beta_1=0.5)\n\n        cycle_gan_compile(monet_generator_optimizer,\n                          photo_generator_optimizer,\n                          monet_discriminator_optimizer,\n                          photo_discriminator_optimizer)\n","metadata":{"execution":{"iopub.status.busy":"2023-02-20T03:10:26.392627Z","iopub.execute_input":"2023-02-20T03:10:26.393074Z","iopub.status.idle":"2023-02-20T03:10:29.959532Z","shell.execute_reply.started":"2023-02-20T03:10:26.393040Z","shell.execute_reply":"2023-02-20T03:10:29.958617Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if __name__ == '__main__':\n    full_dataset = get_gan_dataset_basic(BATCH_SIZE)\n\n    cycle_gan_compile_with_loss_rate(2e-4)\n    cycle_gan_model.fit(full_dataset,\n                        epochs=80)\n\n    cycle_gan_model.monet_generator.save_weights(f'photo2monet.h5')\n    cycle_gan_model.photo_generator.save_weights(f'monet2photo.h5')\n","metadata":{"execution":{"iopub.status.busy":"2023-02-20T03:10:29.960639Z","iopub.execute_input":"2023-02-20T03:10:29.960881Z","iopub.status.idle":"2023-02-20T03:13:15.842670Z","shell.execute_reply.started":"2023-02-20T03:10:29.960849Z","shell.execute_reply":"2023-02-20T03:13:15.841560Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import sys\nfrom pathlib import Path\nfrom numpy import uint8\nfrom PIL import Image\n\nif __name__ == '__main__':\n    photos = photo_dataset().batch(1)\n    try:\n        cycle_gan_model.monet_generator.load_weights(f'photo2monet.h5')\n        cycle_gan_model.photo_generator.load_weights(f'monet2photo.h5')\n    except:\n        sys.exit('Model not trained yet.')\n\n    Path('../submission_images').mkdir(parents=True, exist_ok=True)\n\n    i = 1\n    for img in photos:\n        prediction = cycle_gan_model.monet_generator(img, training=False)[0].numpy()\n        prediction = (prediction * 127.5 + 127.5).astype(uint8)\n        im = Image.fromarray(prediction)\n        im.save(f'../submission_images/{i}.jpg')\n        i += 1\n","metadata":{"execution":{"iopub.status.busy":"2023-02-20T03:13:15.843869Z","iopub.execute_input":"2023-02-20T03:13:15.844128Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import shutil\n\nshutil.make_archive('/kaggle/working/images/', 'zip', '/kaggle/submission_images/')","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}