{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Monet generator\nNOTE: This notebook is WIP - still actively working on finishing it locally. This is a preview.\n\nThe task here is to train a GAN to apply the style of Monet paintings to any input image, or generate Monet-like paintings from scratch. This task is a part of a Kaggle Getting Started competition - more information can be found here: https://www.kaggle.com/competitions/gan-getting-started/data.\nI will train a CycleGAN where all 4 component networks will all use a standard Depthwise Separable Convolutional architecture, but also for the sake of experiment I will test building a version where the convolutional operations are performed using the Fast Fourier Transform algorithm, instead of the standard Keras layers. While using a pre-trained convolutional base for the former could make training faster, for the sake of learning, I will build a UNET-like network from scratch. Large parts of the code have been adapted from the code by Amy Jang, https://www.kaggle.com/code/amyjang/monet-cyclegan-tutorial\n\n***\n## Importing libraries","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport matplotlib.pyplot as plt\nimport PIL\n\n# These imports will be used to pre-process the data\nfrom tensorflow import data\nfrom tensorflow.image import random_brightness, random_flip_left_right, random_flip_up_down, decode_jpeg\nfrom tensorflow.io import gfile, parse_single_example, FixedLenFeature\nfrom tensorflow import string as tf_string\nfrom tensorflow import float32 as tf_float\nfrom tensorflow import reshape as tf_reshape\nfrom tensorflow import cast as tf_cast\n\n# These imports will be used to build, train, and use the network\n#from tensorflow import function as tf_function\nfrom tensorflow import abs as tf_abs\n\nfrom tensorflow import GradientTape, ones_like, zeros_like, reduce_mean\nfrom tensorflow.keras import mixed_precision, Model, Sequential\nfrom tensorflow.keras.initializers import RandomNormal, serialize, RandomUniform, Zeros\nfrom tensorflow.keras.layers import Layer, SeparableConv2D, SpatialDropout2D, LeakyReLU, ReLU, Conv2DTranspose, Input, Concatenate, ZeroPadding2D, Conv2D\nfrom DeepSaki.layers import FourierConvolution2D\nfrom tensorflow_addons.layers import InstanceNormalization\nfrom tensorflow.keras.optimizers import Adam\nfrom tensorflow.keras.losses import BinaryCrossentropy, Reduction\nfrom tensorflow.keras.utils import plot_model\nfrom tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping, Callback\nfrom tensorflow.keras.models import load_model\nfrom os import listdir\nfrom os.path import isfile, normpath\nimport dill\nfrom time import time\n#import keras_tuner","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Defining some network and pre-processing constants","metadata":{}},{"cell_type":"code","source":"BATCH_SIZE = 1\nIMAGE_SIZE = [256, 256]\nCHANNELS = 3\nIMAGE_DIMS = [256, 256, 3]\n\n# Enable float16 mixed precision, to lower the computational complexity (heavily hardware-dependent)\n#mixed_precision.set_global_policy('mixed_float16')","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"***\n## Preparing the data\nThe Monet paintings that will be used to train the discriminator network (the adversary) are saved in a storage-efficient TFRec format. I will first unpack them into tf.data.Dataset objects. Note that I will only be using a few basic image augmentation procedures to artificially increase the size of the dataset - I do not want either of the networks to learn incorrect (augmented) representations of what a Monet painting looks like, style-wise.","metadata":{}},{"cell_type":"code","source":"data_dir = normpath(\"R:\\\\data\")\nmonet = gfile.glob(str(normpath(f'{data_dir}\\\\monet_tfrec\\\\*.tfrec')))\nphotos = gfile.glob(str(normpath(f'{data_dir}\\\\photo_tfrec\\\\*.tfrec')))","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def import_dataset(files, apply_augmentation=False):\n    \"\"\"Extract images from the TFRec files and add them into a tf.data Dataset, optionally applying augmentation to them.\n    Parameters:\n        files (tf.io.gfile.glob object): The list of filenames to be added into the dataset.\n        apply_augmentation (bool): Whether to apply augmentation (random horizontal and vertical flips, random brightness change within a delta of 20%) to copies of the images. Default: False\n    Returns:\n        dataset (tf.data.Datset object): The files as a tf.data Dataset.\"\"\"\n\n    def read_tfrecord(example):\n        \"\"\"Extracts the images from each entry of the TFRec entry, keeping only the images themselves, as the labels are irrelevant.\"\"\"\n        tfrecord_format = {\n            'image_name': FixedLenFeature([], tf_string),\n            'image': FixedLenFeature([], tf_string),\n            'target': FixedLenFeature([], tf_string)\n        }\n\n        example = parse_single_example(example, tfrecord_format)\n        image = decode_jpeg(example['image'], channels=CHANNELS)\n        image = tf_reshape(image, [*IMAGE_SIZE, CHANNELS])\n        return image\n\n\n    ignore_order = data.Options()\n    ignore_order.experimental_deterministic = False\n    dataset = data.TFRecordDataset(files, num_parallel_reads=-1)\n    dataset = dataset.with_options(ignore_order)\n    dataset = dataset.map(read_tfrecord, num_parallel_calls=-1)\n\n    if apply_augmentation:\n        # Random horizontal flips\n        dataset = dataset.concatenate(dataset.map(lambda image: random_flip_left_right(image), num_parallel_calls=-1).shuffle(100000))\n        # Random vertical flips\n        dataset = dataset.concatenate(dataset.map(lambda image: random_flip_up_down(image), num_parallel_calls=-1).shuffle(100000))\n        # Random brightness\n        #dataset = dataset.concatenate(dataset.map(lambda image: random_brightness(image, 0.2), num_parallel_calls=-1).shuffle(100000))\n\n    # Normalize\n    dataset = dataset.map(lambda image: (tf_cast(image, tf_float) / 127.5) - 1, num_parallel_calls=-1)\n\n    return dataset","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"monet = import_dataset(monet, apply_augmentation=True).batch(BATCH_SIZE).cache().prefetch(-1)\nphotos = import_dataset(photos).batch(BATCH_SIZE).cache().prefetch(-1)","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Let's look at an example from each dataset.","metadata":{}},{"cell_type":"code","source":"example_monet = next(iter(monet))\nexample_photo = next(iter(photos))","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.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":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"***\n## Building the GAN\n### Building the AutoEncoder\nBefore building the network itself, I must define an AutoEncoder, which will downsample each image into a lower-dimensional (latent) representation, then upscale it back to the original resolution. I will be using SpatialDropout to apply dropout to the network, and I'll use Depthwise Seprable Convolutions.","metadata":{}},{"cell_type":"code","source":"def downsampler(filters: int, kernel_size: int, stride_count: int, apply_instancenorm: bool, use_dropout: bool, dropout: float, activation: str, use_fft: bool):\n    \"\"\"Defines the downsampling part of the AutoEncoder network, which will be used to reduce the dimensionality of the input images and learn the feature map. Uses depthwise separable convolutions, or FFT-based convolutions.\n    Parameters:\n        filters (int): The number of convolutional filters.\n        kernel_size (int): The size of the convolutional kernel.\n        stride_count (int): The size of the convolutional strides - the length of the convolutional filters.\n        apply_instancenorm (bool): Whether to apply Instance Normalization after convolutions.\n        use_dropout (bool): Whether to apply spatial dropout after convolutions.\n        dropout (float): The dropout rate to be used by the spatial dropout layer. Should be between 0 and 1.\n        activation (string): What type of activation should be applied to the convolutional layers. Should be either 'ReLU' or 'LeakyReLu'.\n        use_fft (bool): Whether to use FFT-based convolutions, or standard depthwise separable convolutions from the TensorFlow library.\n    Returns:\n        x (Keras tensor): The built downsampler.\"\"\"\n\n    initializer = RandomNormal(0., 0.02)\n    gamma_initializer = RandomNormal(mean=0.0, stddev=0.02)\n    x = Sequential()\n\n    if use_fft:\n        x.add(FourierConvolution2D(filters, [kernel_size, kernel_size], use_bias=False, kernel_initializer=initializer, dtype='float32', method='MATRIX_PRODUCT'))\n    else:\n        x.add(SeparableConv2D(filters, kernel_size, strides=stride_count, padding='same', kernel_initializer=initializer, use_bias=False))\n\n    if apply_instancenorm:\n        x.add(InstanceNormalization(gamma_initializer=gamma_initializer))\n\n    if activation == 'LeakyReLU':\n        x.add(LeakyReLU())\n    else:\n        x.add(ReLU())\n\n    if use_dropout:\n        x.add(SpatialDropout2D(dropout))\n\n    return x","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Since TensorFlow does not currently contain any built-in transposed versions of Depthwise Separable convolutions, I have to use standard transposed convolutions, as unofficial implementations I've found were not functional.","metadata":{}},{"cell_type":"code","source":"def upsampler(filters: int, kernel_size: int, stride_count: int, use_dropout: bool, dropout: float, activation: str):\n    \"\"\"Defines the upsampling part of the AutoEncoder network, which will be used to increase the dimensionality of the input images and identify features present in the image. Uses transposed depthwise separable convolutions.\n    Parameters:\n        filters (int): The number of convolutional filters.\n        kernel_size (int): The size of the convolutional kernel.\n        stride_count (int): The size of the convolutional strides - the length of the convolutional filters.\n        use_dropout (bool): Whether to apply spatial dropout after convolutions.\n        dropout (float): The dropout rate to be used by the spatial dropout layer. Should be between 0 and 1.\n        activation (string): What type of activation should be applied to the convolutional layers. Should be either 'ReLU' or 'LeakyReLu'.\n    Returns:\n        x (Keras tensor): The built upsampler.\"\"\"\n\n    initializer = RandomNormal(0., 0.02)\n    gamma_initializer = RandomNormal(mean=0.0, stddev=0.02)\n\n    x = Sequential()\n    x.add(Conv2DTranspose(filters, kernel_size, strides=stride_count, padding='same', kernel_initializer=initializer, use_bias=False))\n\n    x.add(InstanceNormalization(gamma_initializer=gamma_initializer))\n\n    if activation == 'LeakyReLU':\n        x.add(LeakyReLU())\n    else:\n        x.add(ReLU())\n\n    if use_dropout:\n        x.add(SpatialDropout2D(dropout))\n\n    return x","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Building the generator network","metadata":{}},{"cell_type":"code","source":"def Generator(kernel_size: int, stride_count: int, use_fft: bool, apply_instancenorm: bool, use_dropout: bool, dropout_rate: float, activation: str):\n    \"\"\"Defines the generator network.\n    Parameters:\n        kernel_size (int): The size of the convolutional kernel of the AutoEncoder.\n        stride_count (int): The size of the convolutional strides of the AutoEncoder.\n        use_fft (bool): Whether to use FFT-based convolutions, or standard depthwise separable convolutions from the TensorFlow library in the downsamplers.\n        apply_instancenorm (bool): Whether to apply Instance Normalization after convolutions of the AutoEncoder.\n        use_dropout (bool): Whether to apply spatial dropout after convolutions of the AutoEncoder.\n        dropout_rate (float): The dropout rate to be used by the spatial dropout layer of the AutoEncoder. Should be between 0 and 1.\n        activation (string): What type of activation should be applied to the convolutional layers of the AutoEncoder. Should be either 'ReLU' or 'LeakyReLu'.\n    Returns:\n        Model (Keras model): The built generator model.\"\"\"\n\n    if use_fft:\n        stride_count = 1\n    inputs = Input(shape=IMAGE_DIMS, name='generated_image', batch_size=BATCH_SIZE if use_fft else None)\n    x = inputs\n    skips = []\n    dims = [64, 128, 256] + ([512] * 5)\n\n    for i, dim in enumerate(dims):\n        # Applies downsampling to the input 8 times, each with an increasing number of convolutional filters\n        x = downsampler(dim, kernel_size, apply_instancenorm=False if i == 0 else apply_instancenorm, stride_count=stride_count, use_dropout=use_dropout, dropout=dropout_rate, activation='LeakyReLU', use_fft=use_fft)(x)\n        skips.append(x)\n\n    skips = list(reversed(skips[:-1]))\n    dims = list(reversed(dims[:-1]))\n    for i, dim in enumerate(dims):\n        # Applies upsampling to the downsampled input 7 times, each with a decreasing amount of convolutional filters, also symmetrically creating residual connections with the downsampling layers\n        x = upsampler(dim, kernel_size, stride_count, use_dropout, dropout_rate, activation='ReLU')(x)\n        x = Concatenate()([x, skips[i]])\n\n    initializer = RandomNormal(0., 0.02)\n    # Applies one last transposed convolution to the input, with the number of convolutional filters equal to the number of image channels. Note the explicitly-defined dtype - since float16 mixed precision is used, the output needs to be turned back to float32 at the end.\n    output = Conv2DTranspose(CHANNELS, kernel_size, stride_count, kernel_initializer=initializer, dtype='float32', activation='tanh', padding='same')\n    x = output(x)\n\n    return Model(inputs, outputs=x)","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Building the discriminator network","metadata":{}},{"cell_type":"code","source":"def Discriminator(kernel_size: int, stride_count: int, use_fft: bool, apply_instancenorm: bool, use_dropout: bool, dropout_rate: float, activation: str):\n    \"\"\"Defines the discriminator (adversarial) network.\n    Parameters:\n        kernel_size (int): The size of the convolutional kernel of the downsamplers.\n        stride_count (int): The size of the convolutional strides of the downsamplers.\n        use_fft (bool): Whether to use FFT-based convolutions, or standard separable convolutions from the TensorFlow library in the downsamplers.\n        apply_instancenorm (bool): Whether to apply Instance Normalization after convolutions of the downsamplers.\n        use_dropout (bool): Whether to apply spatial dropout after convolutions of the downsamplers.\n        dropout_rate (float): The dropout rate to be used by the spatial dropout layer of the downsamplers. Should be between 0 and 1.\n        activation (string): What type of activation should be applied to the convolutional layers of the downsamplers. Should be either 'ReLU' or 'LeakyReLu'.\n    Returns:\n        Model (Keras model): The built discriminator model.\"\"\"\n\n    initializer = RandomNormal(0., 0.02)\n    gamma_initializer = RandomNormal(mean=0.0, stddev=0.02)\n    x = Sequential()\n    x.add(Input(shape=IMAGE_DIMS, name='input_image', batch_size=BATCH_SIZE if use_fft else None))\n\n    #x = inputs\n    x.add(downsampler(64, kernel_size, stride_count, apply_instancenorm=False, use_dropout=use_dropout, dropout=dropout_rate, activation='LeakyReLU', use_fft=use_fft))\n    x.add(downsampler(128, kernel_size, stride_count, apply_instancenorm, use_dropout, dropout_rate, activation='LeakyReLU', use_fft=use_fft))\n    x.add(downsampler(256, kernel_size, stride_count, apply_instancenorm, use_dropout, dropout_rate, activation='LeakyReLU', use_fft=use_fft))\n    x.add(ZeroPadding2D())\n\n    if use_fft:\n        x.add(FourierConvolution2D(512, [kernel_size, kernel_size], use_bias=False, kernel_initializer=initializer, dtype='float32', method='MATRIX_PRODUCT'))\n    else:\n        x.add(SeparableConv2D(512, kernel_size, 1, kernel_initializer=initializer, use_bias=False, dtype='float32'))\n    x.add(InstanceNormalization(gamma_initializer=gamma_initializer))\n\n    if activation == 'LeakyReLU':\n        x.add(LeakyReLU())\n    else:\n        x.add(ReLU())\n\n    if use_dropout:\n        x.add(SpatialDropout2D(dropout_rate))\n    x.add(ZeroPadding2D())\n\n    # Applies one last convolution to the input, with the number of convolutional filters equal to the number of image channels. Note the explicitly-defined dtype - since float16 mixed precision is used, the output needs to be turned back to float32 at the end.\n    if use_fft:\n        x.add(FourierConvolution2D(1, [kernel_size, kernel_size], kernel_initializer=initializer, dtype='float32', method='MATRIX_PRODUCT'))\n    else:\n        x.add(SeparableConv2D(1, kernel_size, 1, kernel_initializer=initializer, dtype='float32'))\n\n    return x","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Building the GAN model\nNow it's time to combine two copies each of the Generator and the Discriminator into a CycleGAN model.","metadata":{}},{"cell_type":"code","source":"class CycleGAN(Model):\n    def __init__(self, monet_gen, photo_gen, monet_adv, photo_adv, lambda_cycle=10):\n        super(CycleGAN, self).__init__()\n        self.monet_gen = monet_gen\n        self.photo_gen = photo_gen\n        self.monet_adv = monet_adv\n        self.photo_adv = photo_adv\n        self.lambda_cycle = lambda_cycle\n        '''\n        self.monet_gen_loss_tracker = Mean(name='monet_gen_loss')\n        self.monet_adv_loss_tracker = Mean(name='monet_adv_loss')\n        self.photo_gen_loss_tracker = Mean(name='photo_gen_loss')\n        self.photo_adv_loss_tracker = Mean(name='photo_adv_loss')\n\n    @property\n    def metrics(self):\n        return [float(self.monet_gen_loss_tracker.result().numpy()), float(self.monet_adv_loss_tracker.result().numpy()), float(self.photo_gen_loss_tracker.result().numpy()), float(self.photo_adv_loss_tracker.result().numpy())]\n    '''\n\n    def compile(self, gen_loss_fn, adv_loss_fn, cycle_loss_fn, identity_loss_fn, m_gen_opt, p_gen_opt, m_adv_opt, p_adv_opt):\n        super(CycleGAN, self).compile()\n        self.m_gen_opt = m_gen_opt\n        self.p_gen_opt = p_gen_opt\n        self.m_adv_opt = m_adv_opt\n        self.p_adv_opt = p_adv_opt\n        self.gen_loss_fn = gen_loss_fn\n        self.adv_loss_fn = adv_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):\n        real_monet, real_photo = batch\n\n        with GradientTape(persistent=True) as tape:\n            fake_monet = self.monet_gen(real_photo, training=True)\n            cycled_photo = self.photo_gen(fake_monet, training=True)\n\n            fake_photo = self.photo_gen(real_monet, training=True)\n            cycled_monet = self.monet_gen(fake_photo, training=True)\n\n            same_monet = self.monet_gen(real_monet, training=True)\n            same_photo = self.photo_gen(real_photo, training=True)\n\n            adv_real_monet = self.monet_adv(real_monet, training=True)\n            adv_real_photo = self.photo_adv(real_photo, training=True)\n\n            adv_fake_monet = self.monet_adv(fake_monet, training=True)\n            adv_fake_photo = self.photo_adv(fake_photo, training=True)\n\n            monet_gen_loss = self.gen_loss_fn(adv_fake_monet)\n            photo_gen_loss = self.gen_loss_fn(adv_fake_photo)\n\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            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            monet_adv_loss = self.adv_loss_fn(adv_real_monet, adv_fake_monet)\n            photo_adv_loss = self.adv_loss_fn(adv_real_photo, adv_fake_photo)\n\n            monet_gen_grads = tape.gradient(total_monet_gen_loss, self.monet_gen.trainable_variables)\n            photo_gen_grads = tape.gradient(total_photo_gen_loss, self.photo_gen.trainable_variables)\n            monet_adv_grads = tape.gradient(monet_adv_loss, self.monet_adv.trainable_variables)\n            photo_adv_grads = tape.gradient(photo_adv_loss, self.photo_adv.trainable_variables)\n\n            self.m_gen_opt.apply_gradients(zip(monet_gen_grads, self.monet_gen.trainable_variables))\n            self.p_gen_opt.apply_gradients(zip(photo_gen_grads, self.photo_gen.trainable_variables))\n            self.m_adv_opt.apply_gradients(zip(monet_adv_grads,self.monet_adv.trainable_variables))\n            self.p_adv_opt.apply_gradients(zip(photo_adv_grads, self.photo_adv.trainable_variables))\n\n            '''\n            self.monet_gen_loss_tracker.update_state(total_monet_gen_loss)\n            self.monet_adv_loss_tracker.update_state(monet_adv_loss)\n            self.photo_gen_loss_tracker.update_state(total_photo_gen_loss)\n            self.photo_adv_loss_tracker.update_state(photo_adv_loss)\n            '''\n\n            return {\n                'monet_gen_loss': total_monet_gen_loss,\n                'photo_gen_loss': total_photo_gen_loss,\n                'monet_adv_loss': monet_adv_loss,\n                'photo_adv_loss': photo_adv_loss\n            }\n\n    def get_config(self):\n        config = super().get_config()\n        config.update({'monet_gen': self.monet_gen,\n                       'photo_gen': self.photo_gen,\n                       'monet_adv': self.monet_adv,\n                       'photo_adv': self.photo_adv,\n                       'lambda_cycle': self.lambda_cycle})\n        return config","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Loss functions","metadata":{}},{"cell_type":"code","source":"def adv_loss(real, generated):\n    \"\"\"The 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 loss (float).\"\"\"\n    real_loss = BinaryCrossentropy(from_logits=True, reduction=Reduction.NONE)(ones_like(real), real)\n\n    generated_loss = BinaryCrossentropy(from_logits=True, reduction=Reduction.NONE)(zeros_like(generated), generated)\n\n    return (real_loss + generated_loss) * 0.5\n\ndef gen_loss(generated):\n    \"\"\"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.\"\"\"\n    return BinaryCrossentropy(from_logits=True, reduction=Reduction.NONE)(ones_like(generated), generated)\n\ndef cycle_loss(real, cycled, LAMBDA: int):\n    \"\"\"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    return reduce_mean(tf_abs(real - cycled)) * LAMBDA\n\ndef identity_loss(real, same, LAMBDA: int):\n    \"\"\"The 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    return reduce_mean(tf_abs(real - same)) * 0.5 * LAMBDA","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"***\n### Hyperparameter tuning?\nI have originally intended to use Keras Tuner to perform hyperparameter tuning, but after several days of being unable to get it to work using various guides, I've given up. The following few cells are some of the leftovers from these attempts.\n\nI now have to wrap the CycleGAN model in a Hypermodel, which will be passed into a Keras Tuner object to perform hyperparameter tuning. The code here is based on the code by Siraj Hatoum, available here: https://medium.com/@siraj.hatoum/gan-hyperparameter-tuning-with-keras-tuner-81e00ad1d6be","metadata":{}},{"cell_type":"markdown","source":"class HyperGAN(keras_tuner.HyperModel):\n    def __init__(self, use_fft):\n        super(HyperGAN, self).__init__()\n        self.use_fft = use_fft\n\n    def build(self, hp):\n        #hp_kernel_size = hp.Int('kernel_size', 3, 7, 2)\n        #hp_stride_count = hp.Int('stride_count', 2, 6, 2)\n\n        # Define a list of hyperparameters individual to each of the 4 sub-networks\n        hps = []\n        for i in range(4):\n            hps.append([\n                hp.Boolean(f'apply_instancenorm{i}'),\n                hp.Boolean(f'use_dropout{i}'),\n                hp.Float(f'dropout{i}', 0.1, 0.7, 0.2),\n                hp.Choice(f'activation{i}', ['LeakyReLU', 'ReLU'])])\n\n        hp_lr = hp.Float('lr', min_value=1e-6, max_value=1e-3, sampling='log')\n        hp_epsilon = hp.Choice('opt_epsilon', [1.0, 0.1, 1e-7])\n\n        # Define the 2 generators and the 2 adversaries with their respective hyperparameters\n        self.networks = []\n        for i in range(4):\n            if i < 2:\n                self.networks.append(Generator(4, 2, self.use_fft, *hps[i]))\n            else:\n                self.networks.append(Discriminator(4, 2, self.use_fft, *hps[i]))\n\n        model = CycleGAN(*self.networks)\n\n        # Define an individual Adam optimizer for each of the 4 sub-networks\n        optimizers = []\n        for i in range(4):\n            optimizers.append(Adam(learning_rate=hp_lr, epsilon=hp_epsilon))\n\n        model.compile(gen_loss, adv_loss, cycle_loss, identity_loss, *optimizers)\n\n        return model\n\n    def fit(self, hp, model, *args, **kwargs):\n        return model.metrics\n\n    def get_generator(self, model):\n        return model.networks[0]","metadata":{}},{"cell_type":"markdown","source":"tuner = keras_tuner.Hyperband(hypermodel=HyperGAN(use_fft=False), max_epochs=100, directory=normpath('C:/'), objective=[keras_tuner.Objective('monet_gen_loss', 'min'), keras_tuner.Objective('photo_gen_loss', 'min'), keras_tuner.Objective('monet_adv_loss', 'min'), keras_tuner.Objective('photo_adv_loss', 'min')])\ntuner.search_space_summary()","metadata":{}},{"cell_type":"markdown","source":"tuner.search(data.Dataset.zip((monet, photos)), epochs=50, callbacks=[EarlyStopping(patience=5, min_delta=0.0001)])\ntuner.results_summary(1)","metadata":{}},{"cell_type":"markdown","source":"I'll save the best-performing model into the model variable and train it.","metadata":{}},{"cell_type":"markdown","source":"model = HyperGAN(tuner.get_best_hyperparameters()[0])\nmodel.summary()","metadata":{}},{"cell_type":"markdown","source":"***\n### Building the model\nDefining the 2 generators and the 2 adversaries and building the CycleGAN model.","metadata":{}},{"cell_type":"code","source":"hps = [4, 2, False, True, True, 0.5, 'ReLU']\nnetworks = []\nfor i in range(4):\n    if i < 2:\n        networks.append(Generator(*hps))\n    else:\n        networks.append(Discriminator(*hps))\n\n# Define an individual Adam optimizer for each of the 4 sub-networks\noptimizers = []\nfor i in range(4):\n    optimizers.append(Adam(learning_rate=2e-4, epsilon=0.1, decay=1e-6, beta_1=0.5))\n\nmodel = CycleGAN(*networks)\nmodel.compile(gen_loss, adv_loss, cycle_loss, identity_loss, *optimizers)","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Now I'll plot the network structures.","metadata":{}},{"cell_type":"code","source":"plot_model(networks[0], show_shapes=True, show_dtype=True)","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_model(networks[1], show_shapes=True, show_dtype=True)","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_model(networks[2], show_shapes=True, show_dtype=True)","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_model(networks[3], show_shapes=True, show_dtype=True)","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Training the network\nI will first define a custom callback that will keep a (approximate) track of the total training time, and a callback that will log training history to file, then I'll train the network.","metadata":{}},{"cell_type":"code","source":"class TimeTracker(Callback):\n    def __init__(self, filepath, logs=None):\n        self.filepath = filepath\n        if isfile(filepath):\n            with open(filepath, 'rb') as file:\n                self.times = dill.load(file)\n        else:\n            self.times = []\n\n    def on_epoch_begin(self, epoch, logs=None):\n        self.epoch_time_start = time()\n\n    def on_epoch_end(self, epoch, logs=None):\n        self.times.append(time() - self.epoch_time_start)\n        with open(self.filepath, 'wb') as file:\n            dill.dump(self.times, file)","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class HistoryLogger(Callback):\n    def __init__(self, filepath, logs=None):\n        self.filepath = filepath\n        if isfile(filepath):\n            with open(filepath, 'rb') as file:\n                self.history = dill.load(file)\n        else:\n            self.history = {}\n\n    def on_epoch_end(self, epoch, logs=None):\n        if epoch == 0:\n            for key in logs.keys():\n                self.history[key] = [np.mean(logs[key][0])]\n        else:\n            for key in logs.keys():\n                self.history[key].append(np.mean(logs[key][0]))\n        with open(self.filepath, 'wb') as file:\n            dill.dump(self.history, file)","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"results = model.fit(data.Dataset.zip((monet, photos)), epochs=200, callbacks=[ModelCheckpoint(filepath=normpath('E:\\\\cyclegan_checkpoints\\\\checkpoint_{epoch:03d}'), save_weights_only=True),\n                                                                              TimeTracker(normpath('E:\\\\cyclegan_checkpoints\\\\time_history.pkl')),\n                                                                              HistoryLogger(normpath('E:\\\\cyclegan_checkpoints\\\\training_history.pkl'))])","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Continue training from checkpoint if not finished\ndef get_init_epoch(path: str) -> int:\n    file_list = listdir(path)\n    file_list = [int(file.split('_')[-1].split('.')[0]) for file in file_list if file.startswith('checkpoint_')]\n    return max(file_list)\n\ninit_epoch = get_init_epoch(normpath(\"E:\\\\cyclegan_checkpoints\"))\nmodel.load_weights(normpath(f'E:\\\\cyclegan_checkpoints\\\\checkpoint_{init_epoch:03d}'))\nresults = model.fit(data.Dataset.zip((monet, photos)), epochs=200, callbacks=[ModelCheckpoint(filepath=normpath('E:\\\\cyclegan_checkpoints\\\\checkpoint_{epoch:03d}'), save_weights_only=True),\n                                                                              TimeTracker(normpath('E:\\\\cyclegan_checkpoints\\\\time_history.pkl')),\n                                                                              HistoryLogger(normpath('E:\\\\cyclegan_checkpoints\\\\training_history.pkl'))], initial_epoch=init_epoch)","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Results","metadata":{}},{"cell_type":"code","source":"with open(normpath('E:\\\\cyclegan_checkpoints\\\\time_history.pkl'), 'rb') as file:\n    times = dill.load(file)\n    print(f'Total training time: {sum(times) / 3600:.3f} hours')","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with open(normpath('E:\\\\cyclegan_checkpoints\\\\training_history.pkl'), 'rb') as file:\n    results = dill.load(file)","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"epochs = range(len(results['monet_gen_loss']))\nf, axs = plt.subplots(1, 4, figsize=(25,5))\ncolors = ['b', 'r', 'g', 'y']\n\nfor i, metric in enumerate(['monet_gen_loss', 'photo_gen_loss', 'monet_adv_loss', 'photo_adv_loss']):\n    axs[i].plot(epochs, results[metric], colors[i], label=f\"Training {metric}\")\n    axs[i].set_xlabel('Epochs')\n    axs[i].set_ylabel(metric)\n    axs[i].legend()\n    axs[i].grid()\n\nplt.title('Training losses')\nplt.tight_layout()\nplt.show()","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"interpret","metadata":{}},{"cell_type":"markdown","source":"#### Visualization\nLet's look at some examples of Monet-style applied to some images in the test set. I'll first reload the best-performing model weights.","metadata":{}},{"cell_type":"code","source":"model.load_weights(normpath(f'E:\\\\cyclegan_checkpoints\\\\checkpoint_{np.argmin(results[\"monet_gen_loss\"]):03d}'))","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"generator = model.monet_gen\n_, ax = plt.subplots(5, 2, figsize=(12, 12))\nfor i, img in enumerate(photos.take(5)):\n    prediction = 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":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Generating Monet-style images\nI will now use the Monet generator sub-network to apply Monet-style to the images in the test set.","metadata":{}},{"cell_type":"code","source":"i = 1\nfor img in photos:\n    prediction = 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(normpath(f'../images/{i}.jpg'))\n    i += 1","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"***\n## FFT-based model\nI will now try to build a model where the standard convolutional layers will be replaced by FFT-based convolutions. I do not expect any improvements (or worsening) in result quality - rather, I expect a faster training time, due to the fact that FFT-based convolutions have a time complexity of O(n log n), instead of the standard complexity of O(2^n). I will be using a custom layer for this, made by Sascha Kirch, described here: https://towardsdatascience.com/fourier-cnns-with-kernel-sizes-of-1024x1024-and-larger-29f513fd6120","metadata":{}},{"cell_type":"markdown","source":"I will now run hyperparameter tuning again on this new, FFT-based model.","metadata":{}},{"cell_type":"markdown","source":"tuner2 = keras_tuner.Hyperband(hypermodel=HyperGAN(use_fft=True), max_epochs=100, directory=normpath('C:/'))\ntuner2.search_space_summary()","metadata":{}},{"cell_type":"markdown","source":"tuner2.search(data.Dataset.zip((monet, photos)), epochs=50, callbacks=[EarlyStopping(patience=5, min_delta=0.0001)])\ntuner2.results_summary(1)","metadata":{}},{"cell_type":"markdown","source":"model2 = HyperGAN(tuner2.get_best_hyperparameters()[0])\nmodel2.summary()","metadata":{}},{"cell_type":"code","source":"hps[2] = True\nnetworks = []\nfor i in range(4):\n    if i < 2:\n        networks.append(Generator(*hps))\n    else:\n        networks.append(Discriminator(*hps))\n\n# Define an individual Adam optimizer for each of the 4 sub-networks\noptimizers = []\nfor i in range(4):\n    optimizers.append(Adam(learning_rate=2e-4, epsilon=0.1, decay=1e-6, beta_1=0.5))\n\nmodel2 = CycleGAN(*networks)\nmodel2.compile(gen_loss, adv_loss, cycle_loss, identity_loss, *optimizers)","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_model(networks[0], show_shapes=True, show_dtype=True)","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_model(networks[1], show_shapes=True, show_dtype=True)","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_model(networks[2], show_shapes=True, show_dtype=True)","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_model(networks[3], show_shapes=True, show_dtype=True)","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The Python kernel appears to crash when I try to train this below - no error messages. Trying to figure out why.","metadata":{}},{"cell_type":"code","source":"results2 = model2.fit(data.Dataset.zip((monet, photos)), epochs=200, callbacks=[ModelCheckpoint(filepath='.\\\\cyclegan_checkpoints_fft\\\\checkpoint_{epoch:03d}', save_weights_only=True),\n                                                                                TimeTracker(normpath('E:\\\\cyclegan_checkpoints_fft\\\\time_history.pkl')),\n                                                                                HistoryLogger(normpath('E:\\\\cyclegan_checkpoints_fft\\\\training_history.pkl'))])","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Continue training from checkpoint if not finished\ninit_epoch = get_init_epoch(normpath(\"E:\\\\cyclegan_checkpoints_fft\"))\nmodel2.load_weights(normpath(f'E:\\\\cyclegan_checkpoints_fft\\\\checkpoint_{init_epoch:03d}'))\nresults = model2.fit(data.Dataset.zip((monet, photos)), epochs=200, callbacks=[ModelCheckpoint(filepath=normpath('E:\\\\cyclegan_checkpoints_fft\\\\checkpoint_{epoch:03d}'), save_weights_only=True),\n                                                                               TimeTracker(normpath('E:\\\\cyclegan_checkpoints_fft\\\\time_history.pkl')),\n                                                                               HistoryLogger(normpath('E:\\\\cyclegan_checkpoints_fft\\\\training_history.pkl'))], initial_epoch=init_epoch)","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### results","metadata":{}},{"cell_type":"code","source":"with open(normpath('E:\\\\cyclegan_checkpoints_fft\\\\time_history.pkl'), 'rb') as file:\n    times = dill.load(file)\n    print(f'Total training time: {sum(times) / 3600:.3f} hours')","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with open(normpath('E:\\\\cyclegan_checkpoints_fft\\\\training_history.pkl'), 'rb') as file:\n    results2 = dill.load(file)","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"epochs = range(len(results2['monet_gen_loss']))\nf, axs = plt.subplots(1, 2, figsize=(15,5))\n\nfor i, metric in enumerate(['monet_gen_loss', 'photo_gen_loss', 'monet_adv_loss', 'photo_adv_loss']):\n    axs[i].plot(epochs, results2.[metric], colors[i], label=f\"Training {metric}\")\n    axs[i].set_xlabel('Epochs')\n    axs[i].set_ylabel(metric)\n    axs[i].legend()\n    axs[i].grid()\n\nplt.title('Training losses')\nplt.tight_layout()\nplt.show()","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"interpret","metadata":{}},{"cell_type":"code","source":"model2.load_weights(normpath(f'E:\\\\cyclegan_checkpoints_fft\\\\checkpoint_{np.argmin(results2.history[\"monet_gen_loss\"]):03d}'))","metadata":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"generator = model2.monet_gen\n_, ax = plt.subplots(5, 2, figsize=(12, 12))\nfor i, img in enumerate(photos.take(5)):\n    prediction = 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":{"collapsed":false,"jupyter":{"outputs_hidden":false}},"execution_count":null,"outputs":[]}]}