{"cells":[{"metadata":{},"cell_type":"markdown","source":"## In this notebook we are going to implement the new model designed by the creators of CycleGAN, [CUT](https://github.com/taesungp/contrastive-unpaired-translation/) (Contrastive Unpaired Translation), using ResNet as backbone. This model is suposed to be faster to train than CycleGAN getting even better results. As you will be able to see this is not the case for my code so there must be some error somewhere I can't find. If anyone spots it would be nice! You can also check the full code with TF training for easier debug on my github [repository](https://github.com/Brechard/Simple-CUT-TF)\n\n## We also take some function from this [notebook](https://www.kaggle.com/amyjang/monet-cyclegan-tutorial) to load the data.\n"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import tensorflow as tf\nimport tensorflow_addons as tfa\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 datetime\n\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__)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"GCS_PATH = KaggleDatasets().get_gcs_path()\n\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))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Define the functions used to load each photo and apply some data augmentation"},{"metadata":{"trusted":true},"cell_type":"code","source":"IMAGE_SIZE = (256, 256, 3)\nBATCH_SIZE = 1\n\n\ndef decode_image(image, use_augmentation):\n    image = tf.image.decode_jpeg(image, channels=3)\n    if use_augmentation:\n        image = tf.image.random_flip_left_right(image)\n        image = tf.image.random_flip_up_down(image)\n        image = tf.image.random_contrast(image, 0.9, 1.1)\n        image = tf.image.random_brightness(image, 0.1)\n        image = tf.image.random_crop(image, size=IMAGE_SIZE)\n    # Normalize the pixel values in the range [-1, 1]\n    image = (tf.cast(image, tf.float32) / 127.5) - 1\n    image = tf.reshape(image, IMAGE_SIZE)\n    return image\n\n\ndef read_tfrecord(example, use_augmentation):\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'], use_augmentation)\n    return image\n\n\ndef load_dataset(filenames, use_augmentation):\n    dataset = tf.data.TFRecordDataset(filenames)\n    dataset = dataset.map(lambda x: read_tfrecord(x, use_augmentation), num_parallel_calls=AUTOTUNE)\n    dataset = dataset.cache().shuffle(2020).batch(BATCH_SIZE, drop_remainder=True)\n    return dataset.prefetch(tf.data.experimental.AUTOTUNE)\n\ndef display_samples(ds, n_row=4, n_col=7):\n    plt.figure(figsize=(20, int(20 * n_row / n_col)))\n    for i, photo in enumerate(ds.take(n_row * n_col)):\n        plt.subplot(n_row, n_col, i + 1)\n        plt.imshow(photo[0] * 0.5 + 0.5)\n        plt.axis('off')\n    plt.show()\n\nmonet_ds = load_dataset(MONET_FILENAMES, True)\nphoto_ds = load_dataset(PHOTO_FILENAMES, True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"display_samples(photo_ds)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"display_samples(monet_ds)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## We create a common function to plot the history of the trained models"},{"metadata":{"trusted":true},"cell_type":"code","source":"def plot_history_losses(history, suptitle):\n    dict_history = history.history\n    n_cols = 4\n    n_rows = len(dict_history) // n_cols\n    if len(dict_history) % n_cols != 0:\n        n_rows += 1\n    fig = plt.figure(figsize=(20, 10))\n    i = 1\n    for key, items in dict_history.items():\n        ax = fig.add_subplot(n_rows, n_cols, i)\n        ax.plot(items)\n        ax.set_title(key)\n        i += 1\n    plt.suptitle(suptitle)\n    plt.show()\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## The following functions are used to create the ResNet structure. This functions are simple modifications of the code for CycleGAN by [keras] (https://keras.io/examples/generative/cyclegan/).\n\n## The main difference is that we use antialiased Convolutions. Using a modification of he code found [here](https://github.com/adobe/antialiased-cnns/issues/10) ."},{"metadata":{"trusted":true},"cell_type":"code","source":"\nKERNEL_INIT = tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\nGAMMA_INIT = tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n\n\ndef get_filter(filt_size):\n    filter_ = np.array([1., ])\n    if filt_size == 2:\n        filter_ = np.array([1., 1.])\n    elif filt_size == 3:\n        filter_ = np.array([1., 2., 1.])\n    elif filt_size == 4:\n        filter_ = np.array([1., 3., 3., 1.])\n    elif filt_size == 5:\n        filter_ = np.array([1., 4., 6., 4., 1.])\n    elif filt_size == 6:\n        filter_ = np.array([1., 5., 10., 10., 5., 1.])\n    elif filt_size == 7:\n        filter_ = np.array([1., 6., 15., 20., 15., 6., 1.])\n    filter_ = filter_[:, None] * filter_[None, :]\n    filter_ = filter_ / np.sum(filter_)\n    return filter_\n\n\nclass BlurPool(tf.keras.layers.Layer):\n    def __init__(self, filt_size=3, stride=2):\n        super(BlurPool, self).__init__()\n        self.strides = (stride, stride)\n        self.filt_size = filt_size\n\n        self.filter = get_filter(filt_size)\n        self.pad_layer = ReflectionPadding2D()\n\n    def compute_output_shape(self, input_shape):\n        height = input_shape[1] // self.strides[0]\n        width = input_shape[2] // self.strides[1]\n        channels = input_shape[3]\n        return input_shape[0], height, width, channels\n\n    def call(self, x):\n        filter_ = self.filter\n        filter_ = np.tile(filter_[:, :, None, None], (1, 1, tf.keras.backend.int_shape(x)[-1], 1))\n        filter_ = tf.keras.backend.constant(filter_, dtype=tf.keras.backend.floatx())\n        x = self.pad_layer(x)\n        x = tf.keras.backend.depthwise_conv2d(x, filter_, strides=self.strides, padding='valid')\n        return x\n    \n\nclass Upsample(tf.keras.layers.Layer):\n    def __init__(self, filt_size=4, stride=2):\n        super(Upsample, self).__init__()\n        self.filt_size = filt_size\n        self.filt_odd = np.mod(filt_size, 2) == 1\n        self.pad_size = int((filt_size - 1) / 2)\n        self.strides = (stride, stride)\n        self.off = int((stride - 1) / 2.)\n\n        self.filter = get_filter(filt_size=self.filt_size) * (stride ** 2)\n\n    def compute_output_shape(self, input_shape):\n        height = input_shape[1] * self.strides[0]\n        width = input_shape[2] * self.strides[1]\n        channels = input_shape[3]\n        return BATCH_SIZE, height, width, channels\n\n    def call(self, x):\n        filter_ = self.filter\n        filter_ = np.tile(filter_[:, :, None, None], (1, 1, 1, tf.keras.backend.int_shape(x)[-1]))\n        filter_ = tf.keras.backend.constant(filter_, dtype=tf.keras.backend.floatx())\n        ret_val = tf.nn.conv2d_transpose(x, filter_,\n                                         output_shape=self.compute_output_shape(tf.keras.backend.int_shape(x)),\n                                         strides=self.strides)\n        return ret_val\n\n\nclass ReflectionPadding2D(tf.keras.layers.Layer):\n    \"\"\"Implements Reflection Padding as a layer.\n\n    Args:\n        padding(tuple): Amount of padding for the\n        spatial dimensions.\n\n    Returns:\n        A padded tensor with the same type as the input tensor.\n    \"\"\"\n\n    def __init__(self, padding=(1, 1), **kwargs):\n        self.padding = tuple(padding)\n        super(ReflectionPadding2D, self).__init__(**kwargs)\n\n    def call(self, input_tensor, mask=None):\n        padding_width, padding_height = self.padding\n        padding_tensor = [\n            [0, 0],\n            [padding_height, padding_height],\n            [padding_width, padding_width],\n            [0, 0],\n        ]\n        return tf.pad(input_tensor, padding_tensor, mode=\"REFLECT\")\n\n\ndef residual_block(\n        input_tensor,\n        activation,\n        kernel_initializer=KERNEL_INIT,\n        kernel_size=(3, 3),\n        strides=(1, 1),\n        padding=\"valid\",\n        gamma_initializer=GAMMA_INIT,\n        use_bias=False,\n        res_block_n=None\n):\n    dim = input_tensor.shape[-1]\n    input_tensor = layers.Input(input_tensor.shape[1:])\n\n    x = ReflectionPadding2D()(input_tensor)\n    x = layers.Conv2D(\n        dim,\n        kernel_size,\n        strides=strides,\n        kernel_initializer=kernel_initializer,\n        padding=padding,\n        use_bias=use_bias,\n    )(x)\n    x = tfa.layers.InstanceNormalization(gamma_initializer=gamma_initializer)(x)\n    x = activation(x)\n\n    x = ReflectionPadding2D()(x)\n    x = layers.Conv2D(\n        dim,\n        kernel_size,\n        strides=strides,\n        kernel_initializer=kernel_initializer,\n        padding=padding,\n        use_bias=use_bias,\n    )(x)\n    x = tfa.layers.InstanceNormalization(gamma_initializer=gamma_initializer)(x)\n    x = layers.add([input_tensor, x])\n    return tf.keras.models.Model(input_tensor, x, name=f'residual_block_{res_block_n}')\n\n\ndef downsample(\n        input_tensor,\n        filters,\n        activation,\n        kernel_initializer=KERNEL_INIT,\n        kernel_size=(3, 3),\n        padding=\"same\",\n        gamma_initializer=GAMMA_INIT,\n        use_bias=True,\n        norm=True        \n):\n#     TODO: Add no antialias option\n    x = layers.Conv2D(\n        filters,\n        kernel_size,\n        strides=(1, 1),\n        kernel_initializer=kernel_initializer,\n        padding=padding,\n        use_bias=use_bias,\n    )(input_tensor)\n    if norm:\n        x = tfa.layers.InstanceNormalization(gamma_initializer=gamma_initializer)(x)\n    if activation:\n        x = activation(x)\n    x = BlurPool()(x)\n    return x\n\n\ndef upsample(\n        input_tensor,\n        filters,\n        activation,\n        kernel_size=(3, 3),\n        strides=(1, 1),\n        padding=\"same\",\n        kernel_initializer=KERNEL_INIT,\n        gamma_initializer=GAMMA_INIT,\n        use_bias=True\n):\n    x = Upsample()(input_tensor)\n    x = layers.Conv2DTranspose(\n        filters,\n        kernel_size,\n        strides=strides,\n        padding=padding,\n        kernel_initializer=kernel_initializer,\n        use_bias=use_bias,\n    )(x)\n    x = tfa.layers.InstanceNormalization(gamma_initializer=gamma_initializer)(x)\n    if activation:\n        x = activation(x)\n    return x","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## The generator is composed of 2 (+1, input layer) layers. Encoder and Decoder. The difference is important because we only calculate features for the NCE loss from the Encoder."},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_resnet_encoder(filters=64, num_downsampling_blocks=2, num_residual_blocks=4,\n                       gamma_initializer=GAMMA_INIT, name='Encoder'):\n    img_input = layers.Input(shape=IMAGE_SIZE, name=name + \"_img_input\")\n    x = ReflectionPadding2D(padding=(3, 3))(img_input)\n    x = layers.Conv2D(filters, (7, 7), kernel_initializer=KERNEL_INIT, use_bias=False)(x)\n    x = tfa.layers.InstanceNormalization(gamma_initializer=gamma_initializer)(x)\n    x = layers.Activation(\"relu\")(x)\n\n    # Downsampling\n    for i in range(num_downsampling_blocks):\n        filters *= 2\n        x = downsample(x, filters=filters, activation=layers.Activation(\"relu\"))\n\n    # Residual blocks\n    for i in range(num_residual_blocks):\n        x = residual_block(x, activation=layers.Activation(\"relu\"), res_block_n=i)(x)\n\n    return tf.keras.models.Model(img_input, x, name=name)\n\n\ndef get_resnet_decoder(\n        input_shape,\n        filters=64,\n        num_upsample_blocks=2,\n        name='Decoder',\n):\n    img_input = layers.Input(shape=input_shape, name=name + \"_img_input\")\n    x = img_input\n    # Upsampling\n    for i in range(num_upsample_blocks):\n        filters //= 2\n        x = upsample(x, filters, activation=layers.Activation(\"relu\"))\n\n    # Final block\n    x = ReflectionPadding2D(padding=(3, 3))(x)\n    x = layers.Conv2D(3, (7, 7), padding=\"valid\")(x)\n    x = layers.Activation(\"tanh\")(x)\n\n    return tf.keras.models.Model(img_input, x, name=name)\n\n\ndef get_generator(num_downsampling_blocks=2, num_residual_blocks=4, num_upsample_blocks=2):\n    encoder = get_resnet_encoder(num_downsampling_blocks=num_downsampling_blocks, num_residual_blocks=num_residual_blocks)\n    decoder = get_resnet_decoder(encoder.output.shape[1:], num_upsample_blocks=num_upsample_blocks)\n\n    img_input = layers.Input(shape=IMAGE_SIZE, name=\"generator_img_input\")\n    x = img_input\n    x = encoder(x)\n    x = decoder(x)\n    return tf.keras.models.Model(img_input, x, name='Generator')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## The discrimator is the same as the CycleGan in the [keras repository](https://keras.io/examples/generative/cyclegan/)."},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_discriminator(filters=64, kernel_initializer=KERNEL_INIT, num_downsampling=3):\n    img_input = layers.Input(shape=IMAGE_SIZE, name=\"discriminator_img_input\")\n    x = downsample(\n        img_input,\n        filters=filters,\n        activation=layers.LeakyReLU(0.2),\n        kernel_size=(4, 4),\n        norm=False\n    )\n\n    num_filters = filters\n    for num_downsample_block in range(1, num_downsampling):\n        num_filters *= 2\n        x = downsample(\n            x,\n            filters=num_filters,\n            activation=layers.LeakyReLU(0.2),\n            kernel_size=(4, 4),\n        )\n\n    x = layers.Conv2D(num_filters * 2, (4, 4), strides=(1, 1), kernel_initializer=KERNEL_INIT, padding='same', use_bias=True)(x)\n    x = tfa.layers.InstanceNormalization(gamma_initializer=GAMMA_INIT)(x)\n    x = layers.LeakyReLU(0.2)(x)\n\n    x = layers.Conv2D(1, (4, 4), strides=(1, 1), padding=\"same\", kernel_initializer=kernel_initializer)(x)\n\n    model = tf.keras.models.Model(inputs=img_input, outputs=x, name='Discriminator')\n    return model","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## To extract the features of the encoder model we pass the image block by block and save the output in a list that will be returned. The feature layers correspond to the output of the first padding layer, the two 2D conv layers and the first and fifth ResNet blocks."},{"metadata":{"trusted":true},"cell_type":"code","source":"feature_layers = [0, 4, 8, 12, 16]\ndef get_features(encoder, x):\n    \"\"\" Extract the features generated by each block of the encoder model \"\"\"\n    x = encoder.layers[0](x)  # This is just the input layer\n    features = []\n    for i, layer in enumerate(encoder.layers[1:]):\n        x = layer(x)\n        if i in feature_layers:\n            features.append(x)\n    return features\n\n\ndef normalize(x):\n    norm = tf.pow(tf.reduce_sum(tf.math.pow(x, 2), axis=1, keepdims=True), 1 / 2)\n    out = tf.divide(x, norm + 1e-7)\n    return out\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## There is one MLP created for each feature that we will extract."},{"metadata":{"trusted":true},"cell_type":"code","source":"class MLP:\n    def __init__(self, dimension=256, num_patches=64):\n        self.dim = dimension\n        self.num_patches = num_patches\n        self.n_mlps = 0\n\n    def create_mlp(self, feats):\n        for mlp_id, feat in enumerate(feats):\n            mlp = tf.keras.Sequential([tf.keras.layers.Dense(self.dim, input_dim=feat.shape[-1]),\n                                       tf.keras.layers.ReLU(),\n                                       tf.keras.layers.Dense(self.dim)])\n            setattr(self, f'mlp_{mlp_id}', mlp)\n        self.n_mlps = mlp_id + 1\n\n    def forward(self, features, patch_ids=None, use_mlp=True):\n        \"\"\"\n        Forward the features through their corresponding Multi Layer Perceptron.\n        If the Patch IDs are not provided it means that it is the first time being used with this Batch. What we do\n        then is randomly select \"num_patches\" of patches to process and return the IDs so that in the second run we\n        execute the same patches.\n        Args:\n            features: Features extracted from the encoder. Shape must be [Batch size, Heigh, Width, Channels]\n            patch_ids: IDs of the paths to execute. If None, they will be randomly chosen.\n\n        Returns:\n            Processed features and the patch ids.\n        \"\"\"\n        if self.n_mlps == 0 and use_mlp:\n            self.create_mlp(features)\n        return_ids, return_feats = [], []\n        for feat_id, feat in enumerate(features):\n            B, H, W, C = feat.shape\n            feat_reshape = tf.reshape(feat, (B, -1, C))\n\n            if patch_ids is None:\n                patch_id = tf.random.shuffle(tf.range(H * W))[:self.num_patches]\n            else:\n                patch_id = patch_ids[feat_id]\n\n            x_sample = tf.reshape(tf.gather(feat_reshape, patch_id, axis=1), (-1, C))  # reshape(-1, x.shape[1])\n            mlp = getattr(self, f'mlp_{feat_id}')\n            if use_mlp:\n                x_sample = mlp(x_sample)\n            x_sample = normalize(x_sample)\n            return_ids.append(patch_id)\n            return_feats.append(x_sample)\n\n        return return_feats, return_ids","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## The losses. Simple translation to TF from the original [CUT](https://github.com/taesungp/contrastive-unpaired-translation/blob/master/models/patchnce.py) code."},{"metadata":{"trusted":true},"cell_type":"code","source":"def patch_nce_loss(feat_src, feat_tgt, nce_T):\n    n_patches, size = feat_src.shape\n\n    l_pos = tf.matmul(tf.reshape(feat_src, (n_patches, 1, -1)), tf.reshape(feat_tgt, (n_patches, -1, 1)))\n    l_pos = tf.squeeze(l_pos, 1)\n\n    # reshape features to batch size\n\n    l_neg = tf.matmul(feat_src, tf.transpose(feat_tgt))\n\n    # diagonal entries are similarity between same features, and hence meaningless.\n    # Since there is no masked_fill method in tensorflow we will multiply by a unitary matrix with almost zero values in the diagonal\n    diagonal = tf.ones((n_patches, n_patches)) - tf.eye(n_patches) * 0.9999999\n    l_neg = l_neg * diagonal\n\n    out = tf.concat((l_pos, l_neg), axis=1) / nce_T\n    target = [[1] + [0.0] * l_neg.shape[1] for i in range(out.shape[0])]\n#     tf.print(target)\n    loss = tf.keras.losses.CategoricalCrossentropy(reduction=tf.keras.losses.Reduction.NONE, from_logits=True)(target, out)\n#     tf.print('entropy loss:', loss)\n\n    return loss\n\n\ndef gan_loss(y, target_is_real, target_real_label=1.0, target_fake_label=0.0):\n    if target_is_real:\n        return tf.keras.losses.MSE(y, target_real_label)\n\n    return tf.keras.losses.MSE(y, target_fake_label)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Definition of the CUT/FastCUT model."},{"metadata":{"trusted":true},"cell_type":"code","source":"class CUT(tf.keras.Model):\n    def __init__(\n            self,\n            generator,\n            discriminator,\n            mlp,\n            use_mlp=True,\n            lambda_NCE=1.0,\n            lambda_GAN=1.0,\n            fast=False,\n            nce_T=0.07,\n            use_defaults=True\n    ):\n        super(CUT, self).__init__()\n        self.generator = generator\n        self.discriminator = discriminator\n        self.input_layer, self.encoder = generator.layers[:2]\n        self.mlp = mlp\n        self.fast = fast\n        self.use_mlp = True if use_defaults else use_mlp\n        self.lambda_NCE = (10.0 if fast else 1.0) if use_defaults else lambda_NCE\n        self.lambda_GAN = 1.0 if use_defaults else lambda_GAN\n        self.nce_T = 0.07 if use_defaults else nce_T\n\n    def calculate_NCE_loss(self, src, tgt):\n        encoded_features_src = get_features(self.encoder, self.input_layer(src))\n        encoded_features_tgt = get_features(self.encoder, self.input_layer(tgt))\n\n        mlp_features_src, feat_ids = self.mlp.forward(encoded_features_src)\n        mlp_features_tgt, _ = self.mlp.forward(encoded_features_tgt, feat_ids, self.use_mlp)\n\n        total_nce_loss = 0\n        losses = []\n        for feat_src, feat_tgt in zip(mlp_features_src, mlp_features_tgt):\n            nce_loss = self.patch_nce_loss_fn(feat_src, feat_tgt, self.nce_T) * self.lambda_NCE\n            losses.append(tf.reduce_mean(nce_loss))\n            total_nce_loss += tf.reduce_mean(nce_loss)\n        return total_nce_loss / len(encoded_features_tgt), losses\n\n    def compile(\n            self,\n            generator_optimizer,\n            discriminator_optimizer,\n            mlp_optimizer,\n            discriminator_loss_fn,\n            patch_nce_loss_fn\n    ):\n        super(CUT, self).compile()\n        self.generator_optimizer = generator_optimizer\n        self.discriminator_optimizer = discriminator_optimizer\n        self.mlp_optimizer = mlp_optimizer\n        self.discriminator_loss_fn = discriminator_loss_fn\n        self.patch_nce_loss_fn = patch_nce_loss_fn\n        \n    def data_dependent_initialize(self, source_ds, target_ds):\n        \"\"\"\n        The feature network MLP is defined in terms of the shape of the intermediate, extracted\n        features of the encoder portion of the Generator. Because of this, the weights of MLP are\n        initialized at the first feedforward pass with some input images.\n        \"\"\"\n        with tf.GradientTape(persistent=True) as tape:\n            for source, target in zip(source_ds.take(1), target_ds.take(1)):\n                fake_target = self.generator(source, training=True)\n\n                pred_fake = self.discriminator(fake_target, training=True)\n                pred_real = self.discriminator(target, training=True)\n\n                loss_G_GAN = tf.reduce_mean(\n                    self.discriminator_loss_fn(pred_fake, True)) * self.lambda_GAN if self.lambda_GAN > 0 else 0\n\n                total_loss_nce, losses = loss_nce, losses = self.calculate_NCE_loss(source, fake_target) if self.lambda_NCE > 0 else 0\n                loss_nce_identity, losses_identity = 0, [0 for i in losses]\n                if not self.fast and self.lambda_NCE > 0:\n                    fake_identity = self.generator(target, training=True)\n                    loss_nce_identity, losses_identity = self.calculate_NCE_loss(target, fake_identity)\n                    total_loss_nce = (total_loss_nce + loss_nce_identity) / 2\n\n                loss_G = loss_G_GAN + total_loss_nce\n\n                loss_D_fake = tf.reduce_mean(self.discriminator_loss_fn(pred_fake, False))\n                loss_D_real = tf.reduce_mean(self.discriminator_loss_fn(pred_real, True))\n                loss_D = (loss_D_fake + loss_D_real) / 2\n\n                for feat_id, loss in enumerate(losses):\n                    mlp = getattr(self.mlp, f'mlp_{feat_id}')\n                    mlp_gradients = tape.gradient(loss, mlp.trainable_variables)\n                    self.mlp_optimizer.apply_gradients(zip(mlp_gradients, mlp.trainable_variables))\n\n\n    def train_step(self, batch_data):\n        # source is PHOTO and target is MONET\n        source, target = batch_data\n\n        # For FastCUT, we need to calculate different\n        # kinds of losses for the generators and discriminators.\n        # We will perform the following steps here:\n        #\n        # 1. Pass source image through the generator to calculate the fake target image.\n        # 2. Call the discriminator with the fake target and real target.\n        # 3. Calculate the generator loss (adversarial + NCE).\n        # 4. Calculate the discriminator loss.\n        # 5. Update the weights of the generators\n        # 6. Update the weights of the discriminators\n        # 7. Return the losses in a dictionary\n\n        with tf.GradientTape(persistent=True) as tape:\n            # Photo to fake Monet\n            fake_target = self.generator(source, training=True)\n\n            pred_fake = self.discriminator(fake_target, training=True)\n            pred_real = self.discriminator(target, training=True)\n\n            ###########################\n            ##### TRAIN GENERATOR #####\n            ###########################\n\n            # First, G(photo) should fake the discriminator\n            loss_G_GAN = tf.reduce_mean(\n                self.discriminator_loss_fn(pred_fake, True)) * self.lambda_GAN if self.lambda_GAN > 0 else 0\n\n            total_loss_nce, losses = loss_nce, losses = self.calculate_NCE_loss(source, fake_target) if self.lambda_NCE > 0 else 0\n            loss_nce_identity, losses_identity = 0, [0 for i in losses]\n            if not self.fast and self.lambda_NCE > 0:\n                fake_identity = self.generator(target, training=True)\n                loss_nce_identity, losses_identity = self.calculate_NCE_loss(target, fake_identity)\n                total_loss_nce = (total_loss_nce + loss_nce_identity) / 2\n\n            loss_G = loss_G_GAN + total_loss_nce\n\n            ###########################\n            ### TRAIN DISCRIMINATOR ###\n            ###########################\n\n            loss_D_fake = tf.reduce_mean(self.discriminator_loss_fn(pred_fake, False))\n            loss_D_real = tf.reduce_mean(self.discriminator_loss_fn(pred_real, True))\n            loss_D = (loss_D_fake + loss_D_real) / 2\n\n        # Get the gradients for the generators\n        generator_gradients = tape.gradient(loss_G, self.generator.trainable_variables)\n        discriminator_gradients = tape.gradient(loss_D, self.discriminator.trainable_variables)\n        if self.use_mlp:\n            mlp_gradients = []\n            for feat_id, loss in enumerate(losses):\n                mlp = getattr(self.mlp, f'mlp_{feat_id}')\n                mlp_gradients = tape.gradient(loss, mlp.trainable_variables)\n                self.mlp_optimizer.apply_gradients(zip(mlp_gradients, mlp.trainable_variables))\n\n        # Update the weights\n        self.generator_optimizer.apply_gradients(zip(generator_gradients, self.generator.trainable_variables))\n        self.discriminator_optimizer.apply_gradients(zip(discriminator_gradients, self.discriminator.trainable_variables))\n\n        return {\n            \"G_loss\": loss_G,\n            \"D_loss\": loss_D,\n            \"loss_G_GAN\": loss_G_GAN,\n            \"total_loss_nce\": total_loss_nce,\n            \"loss_nce\": loss_nce,\n            \"loss_nce_idt\": loss_nce_identity,\n            'loss_D_fake': loss_D_fake,\n            'loss_D_real': loss_D_real\n        }","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Callback to print images every X amount of epochs"},{"metadata":{"trusted":true},"cell_type":"code","source":"class GANMonitor(tf.keras.callbacks.Callback):\n    \"\"\"A callback to generate and save images after each epoch\"\"\"\n\n    def __init__(self, num_img=2, every_n_epoch=5):\n        self.num_img = num_img\n        self.every_n_epoch = every_n_epoch\n\n    def on_epoch_end(self, epoch, logs=None):\n        if epoch % self.every_n_epoch != 0:\n            return\n        _, ax = plt.subplots(self.num_img, 4, figsize=(20, 10))\n        [ax[0, i].set_title(title) for i, title in enumerate([\"Source\", \"Fake target\", \"Target\", \"Identity target\"])]\n        for i, (source, target) in enumerate(zip(photo_ds.take(self.num_img), monet_ds.take(self.num_img))):\n            prediction = cut_mlp.generator(source)[0].numpy()\n            prediction = (prediction * 127.5 + 127.5).astype(np.uint8)\n            source = (source[0] * 127.5 + 127.5).numpy().astype(np.uint8)\n\n            idt_target = cut_mlp.generator(target)[0].numpy()\n            idt_target = (idt_target * 127.5 + 127.5).astype(np.uint8)\n            target = (target[0] * 127.5 + 127.5).numpy().astype(np.uint8)\n\n            ax[i, 0].imshow(source)\n            ax[i, 1].imshow(prediction)\n            ax[i, 2].imshow(target)\n            ax[i, 3].imshow(idt_target)\n\n            [ax[i, j].axis(\"off\") for j in range(4)]\n            del prediction\n            del source\n            del target\n            del idt_target\n\n\n        plt.show()\n        plt.close()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"nce_T = 0.07\nplotter = GANMonitor()\n\ndef get_model(num_residual_blocks, use_mlp, fast):\n    generator = get_generator(num_residual_blocks=num_residual_blocks)\n    discriminator = get_discriminator()\n    mlp_network = MLP()\n\n    # Create CUT model\n    model = CUT(generator=generator, discriminator=discriminator, mlp=mlp_network, use_mlp=use_mlp, fast=fast)\n\n    # Compile the model\n    model.compile(\n        generator_optimizer=tf.keras.optimizers.Adam(learning_rate=2e-4, beta_1=0.5),\n        discriminator_optimizer=tf.keras.optimizers.Adam(learning_rate=2e-4, beta_1=0.5),\n        mlp_optimizer=tf.keras.optimizers.Adam(learning_rate=2e-4, beta_1=0.5),\n        discriminator_loss_fn=gan_loss,\n        patch_nce_loss_fn=patch_nce_loss\n    )\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"cut_mlp = get_model(num_residual_blocks=9, use_mlp=True, fast=False)\ncut_mlp.data_dependent_initialize(photo_ds, monet_ds)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for layer_id, layer in enumerate(cut_mlp.generator.layers[1].layers[1:]):\n    if layer_id in feature_layers:\n        print(f'Feature layer {layer}')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"history_cut_mlp = cut_mlp.fit(\n    tf.data.Dataset.zip((photo_ds, monet_ds)),\n    epochs=100,\n    callbacks=[plotter],\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"n_images = 2\n_, ax = plt.subplots(n_images, 4, figsize=(20, 10))\n[ax[0, i].set_title(title) for i, title in enumerate([\"Source\", \"Fake target\", \"Target\", \"Identity target\"])]\nfor i, (source, target) in enumerate(zip(photo_ds.take(n_images), monet_ds.take(n_images))):\n    prediction_ = cut_mlp.generator(source)[0].numpy()\n    prediction = (prediction_ * 127.5 + 127.5).astype(np.uint8)\n    source = (source[0] * 127.5 + 127.5).numpy().astype(np.uint8)\n\n    idt_target = cut_mlp.generator(target)[0].numpy()\n    idt_target = (idt_target * 127.5 + 127.5).astype(np.uint8)\n    target = (target[0] * 127.5 + 127.5).numpy().astype(np.uint8)\n\n    ax[i, 0].imshow(source)\n    ax[i, 1].imshow(prediction)\n    ax[i, 2].imshow(target)\n    ax[i, 3].imshow(idt_target)\n    [ax[i, j].axis(\"off\") for j in range(4)]\n\nplt.show()\nplt.close()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import PIL\n! mkdir ../images","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"i = 1\nfor img in photo_ds:\n    prediction = cut_mlp(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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import shutil\nshutil.make_archive(\"/kaggle/working/images\", 'zip', \"/kaggle/images\")\n","execution_count":null,"outputs":[]}],"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":4,"nbformat_minor":4}