{"cells":[{"metadata":{},"cell_type":"markdown","source":"**Import packages and define special layers:**"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"from __future__ import absolute_import, division, print_function\nimport tensorflow as tf\nfrom tensorflow import keras\nfrom tensorflow.keras import layers, optimizers, applications, Sequential\nfrom tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint, LearningRateScheduler, History\nimport tensorflow_addons as tfa\nimport tensorflow_datasets as tfds\n\nfrom kaggle_datasets import KaggleDatasets\nimport glob\nimport imageio\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nimport os\n\nimport sys\nprint(\"Python version:\", sys.version)\nimport PIL\nimport time\nimport shutil\n\nfrom functools import partial\nfrom albumentations import (\n    Compose, RandomBrightness, JpegCompression, HueSaturationValue, RandomContrast, HorizontalFlip,\n    Rotate\n)\n\nfrom IPython import display\n\n\nfrom tensorflow.python.eager import def_function\nfrom tensorflow.python.framework import dtypes\nfrom tensorflow.python.framework import tensor_shape\nfrom tensorflow.python.keras import backend as K\nfrom tensorflow.python.keras import layers\nfrom tensorflow.python.keras import initializers\nfrom tensorflow.python.ops import array_ops\nfrom tensorflow.python.ops import math_ops\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#Define minibatch discrimination layer\n# https://github.com/garridoq/gan-guide/blob/master/Minibatch%20discrimination%20%2B%20label%20smoothing.ipynb\nclass MinibatchDiscrimination(tf.keras.layers.Layer):\n\n    def __init__(self, num_kernel, dim_kernel,kernel_initializer='glorot_uniform', **kwargs):\n        self.num_kernel = num_kernel\n        self.dim_kernel = dim_kernel\n        self.kernel_initializer = kernel_initializer\n        super(MinibatchDiscrimination, self).__init__(**kwargs)\n\n    def build(self, input_shape):\n        # Create a trainable weight variable for this layer.\n        self.kernel = self.add_weight(name='kernel', \n                                      shape=(input_shape[-1], self.num_kernel*self.dim_kernel),\n                                      initializer=self.kernel_initializer,\n                                      trainable=True)\n        super(MinibatchDiscrimination, self).build(input_shape)\n        \n    def call(self, x):\n        activation = tf.matmul(x, self.kernel)\n        activation = tf.reshape(activation, shape=(-1, self.num_kernel, self.dim_kernel))\n        #Mi\n        tmp1 = tf.expand_dims(activation, 3)\n        #Mj\n        tmp2 = tf.transpose(activation, perm=[1, 2, 0])\n        tmp2 = tf.expand_dims(tmp2, 0)\n        \n        diff = tmp1 - tmp2\n        \n        l1 = tf.reduce_sum(tf.math.abs(diff), axis=2)\n        features = tf.reduce_sum(tf.math.exp(-l1), axis=2)\n        return tf.concat([x, features], axis=1)\n    \n    def compute_output_shape(self, input_shape):\n        return (input_shape[0], input_shape[1] + self.num_kernel)\n\n#Define spectral normalization layer\n# https://medium.com/@FloydHsiu0618/spectral-normalization-implementation-of-tensorflow-2-0-keras-api-d9060d26de77\nclass SpectralNormalization(tf.keras.layers.Layer): # tried layers.Wrapper, also got error\n    def __init__(self, layer, iteration=1, **kwargs):\n        super(SpectralNormalization, self).__init__(layer, **kwargs)\n        self.layer = layer\n        self.iteration = iteration\n\n    def build(self):\n\n        if not self.layer.built:\n            self.layer.build(input_shape)\n\n            if not hasattr(self.layer, 'kernel'):\n                raise ValueError('Invalid layer for SpectralNormalization.')\n\n            self.w = self.layer.kernel\n            self.w_shape = self.w.shape.as_list()\n            self.u = self.add_variable(shape=(1, self.w_shape[-1]), initializer=tf.random_normal_initializer(), name='sn_u', trainable=False, dtype=tf.float32)\n\n        super(SpectralNormalization, self).build()\n\n    @tf.function\n    def call(self, inputs, training=None):\n\n        self._compute_weights(training)\n        output = self.layer(inputs)\n\n        return output\n\n    def _compute_weights(self, training):\n       \n        iteration = self.iteration\n        w_reshaped = tf.reshape(self.w, [-1, self.w_shape[-1]])\n\n        u_hat = tf.identity(self.u)\n        v_hat = None\n\n        for _ in range(self.iteration):\n               \n            v_ = tf.matmul(u_hat, tf.transpose(w_reshaped))\n            v_hat = tf.nn.l2_normalize(v_)\n\n            u_ = tf.matmul(v_hat, w_reshaped)\n            u_hat = tf.nn.l2_normalize(u_)\n\n        if training == True: self.u.assign(u_hat)\n        sigma = tf.matmul(tf.matmul(v_hat, w_reshaped), tf.transpose(u_hat))\n       \n        w_norm = self.w / sigma\n\n        self.layer.kernel = w_norm\n       \n    def compute_output_shape(self, input_shape):\n\n        return tf.TensorShape(self.layer.compute_output_shape(input_shape).as_list())","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"***Set the accelerator to TPU* and then run the following code:**"},{"metadata":{"trusted":true},"cell_type":"code","source":"try:\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":{},"cell_type":"markdown","source":"**Load the data:**\n\n(The various definitions are used to augment the small data set)"},{"metadata":{"trusted":true},"cell_type":"code","source":"from kaggle_datasets import KaggleDatasets\nGCS_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\nIMAGE_SIZE = [256, 256]\n\ndef normalize(image):\n    return (tf.cast(image, tf.float32) / 127.5) - 1\n\ndef decode_image(image):\n    image = tf.image.decode_jpeg(image, channels=3)\n    image = tf.reshape(image, [*IMAGE_SIZE, 3])\n    return image\n\ndef random_crop(image):\n    cropped_image = tf.image.random_crop(image, size=[256, 256, 3])\n    return cropped_image\n\ndef random_jitter(image):\n    # resizing to 286 x 286 x 3 \n    image = tf.image.resize(image, [int(256*1.3), int(256*1.3)],\n                          method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)\n    # randomly cropping to 256 x 256 x 3\n    image = random_crop(image)\n    # random mirroring\n    return image\n\ndef flip(image):\n    return tf.image.flip_left_right(image)\n\ndef preprocess_image_train(image, label=None):\n    image = random_jitter(image)\n    return image\n\ndef read_tfrecord(example):\n    tfrecord_format = {\n        \"image_name\": tf.io.FixedLenFeature([], tf.string),\n        \"image\": tf.io.FixedLenFeature([], tf.string),\n        \"target\": tf.io.FixedLenFeature([], tf.string)\n    }\n    example = tf.io.parse_single_example(example, tfrecord_format)\n    image = decode_image(example['image'])\n    return image\n\ndef load_dataset(filenames, labeled=False, ordered=False, repeats=200):\n    dataset = tf.data.TFRecordDataset(filenames)\n    dataset = dataset.map(read_tfrecord, num_parallel_calls=AUTOTUNE)\n    dataset = dataset.concatenate(dataset.map(flip, num_parallel_calls=AUTOTUNE).shuffle(100000))\n    dataset = dataset.concatenate(dataset.map(random_jitter, num_parallel_calls=AUTOTUNE).shuffle(10000, reshuffle_each_iteration=True).repeat(repeats))\n    dataset = dataset.map(normalize, num_parallel_calls=AUTOTUNE).shuffle(10000)\n    return dataset\n\nmonet_ds = load_dataset(MONET_FILENAMES, labeled=True, repeats=100).batch(100, drop_remainder=True)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Display the first images from the dataset:**"},{"metadata":{"trusted":true},"cell_type":"code","source":"def view_image(ds, rows=2):\n    image = next(iter(ds)) # extract 1 batch from the dataset\n    image = image.numpy()\n\n    fig = plt.figure(figsize=(22, rows * 5.05 ))\n    for i in range(5 * rows):\n        ax = fig.add_subplot(rows, 5, i+1, xticks=[], yticks=[])\n        ax.imshow(image[i] / 2 + .5)\n\nview_image(monet_ds)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Create the generator:**\n\n(The generator samples noise, reshapes and upsamples it to Monetize according to the current distribution. Choose the appropriate architecture for experiments)"},{"metadata":{"trusted":true},"cell_type":"code","source":"#Original generator architecture with batch normalization and LeakyReLUs\ndef Generator(LATENT_DIM=128, OUTPUT_CHANNELS=3):\n    model = tf.keras.Sequential()\n    model.add(layers.Dense(4*4*LATENT_DIM, use_bias=False, input_shape=(LATENT_DIM,)))\n    model.add(layers.BatchNormalization())\n    model.add(layers.LeakyReLU(alpha=0.2))\n\n    model.add(layers.Reshape((4, 4, LATENT_DIM)))\n    #assert model.output_shape == (None, 4, 4, LATENT_DIM) # Note: None is the batch size\n\n    initializer = tf.random_normal_initializer(0., 0.02)\n    model.add(layers.Conv2DTranspose(LATENT_DIM, 4, strides=(2, 2), padding='same', kernel_initializer=initializer, use_bias=False))\n    #assert model.output_shape == (None, 8, 8, LATENT_DIM)\n    model.add(layers.BatchNormalization())\n    model.add(layers.LeakyReLU())\n\n    initializer = tf.random_normal_initializer(0., 0.02)\n    model.add(layers.Conv2DTranspose(LATENT_DIM, 4, strides=(2, 2), padding='same', kernel_initializer=initializer, use_bias=False))\n    #assert model.output_shape == (None, 16, 16, LATENT_DIM)\n    model.add(layers.BatchNormalization())\n    model.add(layers.LeakyReLU())\n    \n    initializer = tf.random_normal_initializer(0., 0.02)\n    model.add(layers.Conv2DTranspose(LATENT_DIM//2, 4, strides=(2, 2), padding='same', kernel_initializer=initializer, use_bias=False))\n    #assert model.output_shape == (None, 32, 32, LATENT_DIM//2)\n    model.add(layers.BatchNormalization())\n    model.add(layers.LeakyReLU())\n    \n    initializer = tf.random_normal_initializer(0., 0.02)\n    model.add(layers.Conv2DTranspose(LATENT_DIM//4, 4, strides=(2, 2), padding='same', kernel_initializer=initializer, use_bias=False))\n    #assert model.output_shape == (None, 64, 64, LATENT_DIM//4)\n    model.add(layers.BatchNormalization())\n    model.add(layers.LeakyReLU())\n\n    initializer = tf.random_normal_initializer(0., 0.02)\n    model.add(layers.Conv2DTranspose(LATENT_DIM//8, 4, strides=(2, 2), padding='same', kernel_initializer=initializer, use_bias=False))\n    #assert model.output_shape == (None, 128, 128, LATENT_DIM//8)\n    model.add(layers.BatchNormalization())\n    model.add(layers.LeakyReLU())\n    \n    initializer = tf.random_normal_initializer(0., 0.02)\n    model.add(layers.Conv2DTranspose(OUTPUT_CHANNELS, 4, strides=(2, 2), padding='same', kernel_initializer=initializer, use_bias=False, activation='tanh'))\n    #assert model.output_shape == (None, 256, 256, 4)\n    #model.add(layers.BatchNormalization())\n    #model.add(layers.LeakyReLU())\n\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#Generator with spectral normalization and LeakyReLUs\ndef Generator_s(LATENT_DIM=128, OUTPUT_CHANNELS=3):\n    model = tf.keras.Sequential()\n    model.add(SpectralNormalization(layers.Dense(4*4*LATENT_DIM, use_bias=False, input_shape=(LATENT_DIM,))))\n    #model.add(SpectralNormalization())\n    model.add(layers.LeakyReLU(alpha=0.2))\n\n    model.add(layers.Reshape((4, 4, LATENT_DIM)))\n    #assert model.output_shape == (None, 4, 4, LATENT_DIM) # Note: None is the batch size\n\n    initializer = tf.random_normal_initializer(0., 0.02)\n    model.add(SpectralNormalization(layers.Conv2DTranspose(LATENT_DIM, 4, strides=(2, 2), padding='same', kernel_initializer=initializer, use_bias=False)))\n    #assert model.output_shape == (None, 8, 8, LATENT_DIM)\n    #model.add(SpectralNormalization())\n    model.add(layers.LeakyReLU())\n\n    initializer = tf.random_normal_initializer(0., 0.02)\n    model.add(SpectralNormalization(layers.Conv2DTranspose(LATENT_DIM, 4, strides=(2, 2), padding='same', kernel_initializer=initializer, use_bias=False)))\n    #assert model.output_shape == (None, 16, 16, LATENT_DIM)\n    #model.add(SpectralNormalization())\n    model.add(layers.LeakyReLU())\n    \n    initializer = tf.random_normal_initializer(0., 0.02)\n    model.add(SpectralNormalization(layers.Conv2DTranspose(LATENT_DIM//2, 4, strides=(2, 2), padding='same', kernel_initializer=initializer, use_bias=False)))\n    #assert model.output_shape == (None, 32, 32, LATENT_DIM//2)\n    #model.add(SpectralNormalization())\n    model.add(layers.LeakyReLU())\n    \n    initializer = tf.random_normal_initializer(0., 0.02)\n    model.add(SpectralNormalization(layers.Conv2DTranspose(LATENT_DIM//4, 4, strides=(2, 2), padding='same', kernel_initializer=initializer, use_bias=False)))\n    #assert model.output_shape == (None, 64, 64, LATENT_DIM//4)\n    #model.add(SpectralNormalization())\n    model.add(layers.LeakyReLU())\n\n    initializer = tf.random_normal_initializer(0., 0.02)\n    model.add(SpectralNormalization(layers.Conv2DTranspose(LATENT_DIM//8, 4, strides=(2, 2), padding='same', kernel_initializer=initializer, use_bias=False)))\n    #assert model.output_shape == (None, 128, 128, LATENT_DIM//8)\n    #model.add(SpectralNormalization())\n    model.add(layers.LeakyReLU())\n    \n    initializer = tf.random_normal_initializer(0., 0.02)\n    model.add(SpectralNormalization(layers.Conv2DTranspose(OUTPUT_CHANNELS, 4, strides=(2, 2), padding='same', kernel_initializer=initializer, use_bias=False, activation='tanh')))\n    #assert model.output_shape == (None, 256, 256, 4)\n    #model.add(SpectralNormalization())\n    model.add(layers.LeakyReLU())\n\n    return model","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Create the discriminator:**\n\n(The discriminator takes in the input image and classifies it as real or fake (generated). But instead of outputting a single node, the discriminator outputs a smaller 2D image with higher pixel values indicating a real classification and lower values indicating a fake classification. Choose the appropriate architecture for experiments)"},{"metadata":{"trusted":true},"cell_type":"code","source":"#Original discriminator with LeakyReLUs and dropout=0.3, final layer has sigmoid activation function\ndef Discriminator():\n    model = tf.keras.Sequential()\n    initializer = tf.random_normal_initializer(0., 0.02)\n    model.add(layers.Conv2D(64, 4, strides=(2, 2), padding='same', kernel_initializer=initializer,\n                                     input_shape=[256, 256, 3], use_bias=False))\n    model.add(layers.LeakyReLU())\n    model.add(layers.Dropout(0.1))\n\n    initializer = tf.random_normal_initializer(0., 0.02)\n    model.add(layers.Conv2D(128, 4, strides=(2, 2), padding='same', kernel_initializer=initializer, use_bias=False))\n    model.add(layers.LeakyReLU())\n    model.add(layers.Dropout(0.1))\n    \n    initializer = tf.random_normal_initializer(0., 0.02)\n    model.add(layers.Conv2D(256, 4, strides=(2, 2), padding='same', kernel_initializer=initializer, use_bias=False))\n    model.add(layers.LeakyReLU())\n    model.add(layers.Dropout(0.1))\n    \n    model.add(layers.ZeroPadding2D())\n    initializer = tf.random_normal_initializer(0., 0.02)\n    gamma_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n    model.add(layers.Conv2D(512, 4, strides=1, kernel_initializer=initializer, use_bias=False))\n    model.add(tfa.layers.InstanceNormalization(gamma_initializer=gamma_init))\n    model.add(layers.LeakyReLU())\n    model.add(layers.ZeroPadding2D())\n    model.add(layers.Conv2D(1, 4, strides=1, kernel_initializer=initializer))\n    model.add(layers.LeakyReLU(alpha=0.2))\n    \n    model.add(layers.Flatten())\n    model.add(layers.Dense(1, activation='linear'))\n\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#Discriminator with addition of minibatch discrimination.\ndef Discriminator_m():\n    model = tf.keras.Sequential()\n    initializer = tf.random_normal_initializer(0., 0.02)\n    model.add(layers.Conv2D(64, 4, strides=(2, 2), padding='same', kernel_initializer=initializer,\n                                     input_shape=[256, 256, 3], use_bias=False))\n    model.add(layers.LeakyReLU())\n    model.add(layers.Dropout(0.3))\n\n    initializer = tf.random_normal_initializer(0., 0.02)\n    model.add(layers.Conv2D(128, 4, strides=(2, 2), padding='same', kernel_initializer=initializer, use_bias=False))\n    model.add(layers.LeakyReLU())\n    model.add(layers.Dropout(0.3))\n    \n    initializer = tf.random_normal_initializer(0., 0.02)\n    model.add(layers.Conv2D(256, 4, strides=(2, 2), padding='same', kernel_initializer=initializer, use_bias=False))\n    model.add(layers.LeakyReLU())\n    model.add(layers.Dropout(0.3))\n    \n    model.add(layers.ZeroPadding2D())\n    initializer = tf.random_normal_initializer(0., 0.02)\n    gamma_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n    model.add(layers.Conv2D(512, 4, strides=1, kernel_initializer=initializer, use_bias=False))\n    model.add(tfa.layers.InstanceNormalization(gamma_initializer=gamma_init))\n    model.add(layers.LeakyReLU())\n    model.add(layers.BatchNormalization())\n    model.add(layers.Flatten())\n    model.add(layers.Dense(512))\n    model.add(layers.LeakyReLU())\n    model.add(MinibatchDiscrimination(num_kernel=100, dim_kernel=5))\n    model.add(layers.Dense(512))\n    model.add(layers.LeakyReLU(alpha=0.2))\n    model.add(layers.Dense(1, activation='linear'))\n\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#Discriminator with spectral normalization\ndef Discriminator_s():\n    model = tf.keras.Sequential()\n    initializer = tf.random_normal_initializer(0., 0.02)\n    model.add(layers.Conv2D(64, 4, strides=(2, 2), padding='same', kernel_initializer=initializer,\n                                     input_shape=[256, 256, 3], use_bias=False))\n    model.add(SpectralNormalization())\n    model.add(layers.LeakyReLU())\n    model.add(layers.Dropout(0.3))\n\n    initializer = tf.random_normal_initializer(0., 0.02)\n    model.add(layers.Conv2D(128, 4, strides=(2, 2), padding='same', kernel_initializer=initializer, use_bias=False))\n    model.add(SpectralNormalization())\n    model.add(layers.LeakyReLU())\n    model.add(layers.Dropout(0.3))\n    \n    initializer = tf.random_normal_initializer(0., 0.02)\n    model.add(layers.Conv2D(256, 4, strides=(2, 2), padding='same', kernel_initializer=initializer, use_bias=False))\n    model.add(SpectralNormalization())\n    model.add(layers.LeakyReLU())\n    model.add(layers.Dropout(0.3))\n    \n    model.add(layers.ZeroPadding2D())\n    initializer = tf.random_normal_initializer(0., 0.02)\n    gamma_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n    model.add(layers.Conv2D(512, 4, strides=1, kernel_initializer=initializer, use_bias=False))\n    model.add(SpectralNormalization())\n    model.add(layers.LeakyReLU())\n    model.add(layers.ZeroPadding2D())\n    model.add(layers.Conv2D(1, 4, strides=1, kernel_initializer=initializer))\n    model.add(layers.LeakyReLU(alpha=0.2))\n    \n    model.add(layers.Flatten())\n    model.add(layers.Dense(1, activation='sigmoid'))\n\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#Discriminator with minibatch discrimination and spectral normalization.\ndef Discriminator_ms():\n    model = tf.keras.Sequential()\n    initializer = tf.random_normal_initializer(0., 0.02)\n    model.add(SpectralNormalization(layers.Conv2D(64, 4, strides=(2, 2), padding='same', kernel_initializer=initializer, input_shape=[256, 256, 3], use_bias=False)))\n    model.add(layers.LeakyReLU())\n    model.add(layers.Dropout(0.3))\n\n    initializer = tf.random_normal_initializer(0., 0.02)\n    model.add(SpectralNormalization(layers.Conv2D(128, 4, strides=(2, 2), padding='same', kernel_initializer=initializer, use_bias=False)))\n    model.add(layers.LeakyReLU())\n    model.add(layers.Dropout(0.3))\n    \n    initializer = tf.random_normal_initializer(0., 0.02)\n    model.add(SpectralNormalization(layers.Conv2D(256, 4, strides=(2, 2), padding='same', kernel_initializer=initializer, use_bias=False)))\n    model.add(layers.LeakyReLU())\n    model.add(layers.Dropout(0.3))\n    \n    model.add(layers.ZeroPadding2D())\n    initializer = tf.random_normal_initializer(0., 0.02)\n    gamma_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n    model.add(SpectralNormalization(layers.Conv2D(512, 4, strides=1, kernel_initializer=initializer, use_bias=False)))\n    model.add(layers.LeakyReLU())\n    model.add(layers.BatchNormalization())\n    model.add(layers.Flatten())\n    model.add(layers.Dense(512))\n    model.add(layers.LeakyReLU())\n    model.add(MinibatchDiscrimination(num_kernel=100, dim_kernel=5))\n    model.add(layers.Dense(512))\n    model.add(layers.LeakyReLU(alpha=0.2))\n    model.add(layers.Dense(1, activation='sigmoid'))\n\n    return model","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Define loss:**"},{"metadata":{"trusted":true},"cell_type":"code","source":"#Least squares loss\nwith strategy.scope():\n    def discriminator_loss(predictions_real, predictions_gen, labels_real):\n        gen_loss  = tf.reduce_mean((predictions_gen  - tf.reduce_mean(predictions_real) + labels_real) ** 2)\n        real_loss = tf.reduce_mean((predictions_real - tf.reduce_mean(predictions_gen)  - labels_real) ** 2)\n        return (gen_loss + real_loss) / 2\n    \n    def generator_loss(predictions_real, predictions_gen, labels_real):\n        gen_loss  = tf.reduce_mean((predictions_gen  - tf.reduce_mean(predictions_real) - labels_real) ** 2)\n        real_loss = tf.reduce_mean((predictions_real - tf.reduce_mean(predictions_gen)  + labels_real) ** 2)\n        return (gen_loss + real_loss) / 2","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Define GAN monitor (for images)**"},{"metadata":{"trusted":true},"cell_type":"code","source":"gen_dir = 'generated'\n\nif not os.path.exists(gen_dir):\n    os.makedirs(gen_dir)\n\nclass GANMonitor(keras.callbacks.Callback):\n    def __init__(self, num_img=3, latent_dim=128):\n        self.num_img = num_img\n        self.latent_dim = latent_dim\n\n    def on_epoch_end(self, epoch, logs=None):\n        random_latent_vectors = tf.random.normal(shape=(self.num_img, self.latent_dim))\n        generated_images = self.model.generator(random_latent_vectors)\n        generated_images *= 255\n        generated_images.numpy()\n        for i in range(self.num_img):\n            img = keras.preprocessing.image.array_to_img(generated_images[i])\n            img.save(\"generated/generated_img_{i}_{epoch}.png\".format(i=i, epoch=epoch))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Define MonetGAN**"},{"metadata":{"trusted":true},"cell_type":"code","source":"OUTPUT_CHANNELS = 3\nLATENT_DIM = 128\n\nwith strategy.scope():\n    monet_generator = Generator(LATENT_DIM, 3)\n    monet_discriminator = Discriminator()\n\nclass MonetGan(keras.Model):\n    def __init__(self, monet_generator, monet_discriminator, latent_dim, real_label=0.5, fake_label=0):\n        super(MonetGan, self).__init__()\n        self.generator = monet_generator\n        self.discriminator = monet_discriminator\n        self.latent_dim = latent_dim\n        self.real_label = real_label\n        self.fake_label = fake_label\n        \n    def compile(self, d_opt, g_opt, d_loss_fn, g_loss_fn):\n        super(MonetGan, self).compile()\n        self.d_opt = d_opt\n        self.g_opt = g_opt\n        self.d_loss_fn = d_loss_fn\n        self.g_loss_fn = g_loss_fn\n        \n    def train_step(self, images):\n        #noise = tf.random.normal([BATCH_SIZE, noise_dim])\n    \n        if isinstance(images, tuple):\n            images = images[0]\n    \n        # Sample random points in the latent space\n        batch_size = tf.shape(images)[0]\n        noise = tf.random.normal(shape=(batch_size, self.latent_dim))\n        \n        #labels_gen  = tf.zeros((batch_size, 1)) + fake_label\n        labels_real = tf.zeros((batch_size, 1)) + self.real_label\n        \n        # Add random noise to the labels - important trick!\n        #labels_gen  += 0.05 * tf.random.uniform(tf.shape(labels_gen))\n        labels_real += 0.05 * tf.random.uniform(tf.shape(labels_real))\n    \n        with tf.GradientTape() as disc_tape: \n            generated_images = self.generator(noise, training=False)\n        \n            real_output = self.discriminator(images, training=True)\n            fake_output = self.discriminator(generated_images, training=True)\n    \n            disc_loss = self.d_loss_fn(real_output, fake_output, labels_real)\n    \n        with tf.GradientTape() as gen_tape: \n            generated_images = self.generator(noise, training=True)\n        \n            real_output = self.discriminator(images, training=False)\n            fake_output = self.discriminator(generated_images, training=False)\n    \n            gen_loss = self.g_loss_fn(real_output, fake_output, labels_real)\n\n        gradients_of_discriminator = disc_tape.gradient(disc_loss, self.discriminator.trainable_variables)\n        gradients_of_generator = gen_tape.gradient(gen_loss, self.generator.trainable_variables)\n\n        self.g_opt.apply_gradients(zip(gradients_of_generator, self.generator.trainable_variables))\n        self.d_opt.apply_gradients(zip(gradients_of_discriminator, self.discriminator.trainable_variables))\n    \n        return {\"d_loss\": disc_loss, \"g_loss\": gen_loss}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Define evaluation metric"},{"metadata":{},"cell_type":"markdown","source":"**Training parameters and checkpoints (creates new model)**"},{"metadata":{"trusted":true},"cell_type":"code","source":"checkpoint_path = 'training_1/cp-{epoch:04d}.h5'\ncheckpoint_dir = os.path.dirname(checkpoint_path)\nif not os.path.exists(checkpoint_dir):\n    os.makedirs(checkpoint_dir)\n\n\n# Creates checkpoint callback to pass to model.fit\ncheckpoint = ModelCheckpoint(checkpoint_path,\n                            save_weights_only=True,\n                            verbose=1,\n                            periods=5)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#Training \nEPOCHS = 50\n\nLR_G = 0.001\nLR_D = 0.0005\nbeta_1 = .5\n\nreal_label = 0.05\nfake_label = 0.95\n\nwith strategy.scope():\n    monet_gan = MonetGan(monet_discriminator=monet_discriminator, \n                         monet_generator=monet_generator, \n                         latent_dim=LATENT_DIM,\n                         real_label=real_label,\n                         fake_label=fake_label)\n    \n    monet_gan.compile(\n        d_opt = tf.keras.optimizers.Adam(learning_rate=LR_D, beta_1=beta_1),\n        g_opt = tf.keras.optimizers.Adam(learning_rate=LR_G, beta_1=beta_1),\n        d_loss_fn=discriminator_loss,\n        g_loss_fn=generator_loss\n    )","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Look at saved checkpoints (if not training from scratch)**"},{"metadata":{"trusted":true},"cell_type":"code","source":"!ls {checkpoint_dir}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Choose and load latest checkpoint (if not training from scratch)**"},{"metadata":{"trusted":true},"cell_type":"code","source":"list_of_files = glob.glob(checkpoint_dir + '/*.h5') # * means all if need specific format then *.csv\nlatest_file = max(list_of_files, key=os.path.getctime)\nprint(latest_file)\n# Load checkpoint\nmonet_gan.load_weights(latest_file)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Only run this in case you wanna delete folders"},{"metadata":{"trusted":true},"cell_type":"code","source":"#shutil.rmtree(\"/kaggle/working/generated\")\n#shutil.rmtree(\"/kaggle/working/training_1\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"[Learning rate scheduler](https://keras.io/guides/writing_your_own_callbacks/#learning-rate-scheduling)"},{"metadata":{"trusted":true},"cell_type":"code","source":"class CustomLearningRateScheduler(keras.callbacks.Callback):\n    \"\"\"Learning rate scheduler which sets the learning rate according to schedule.\n\n  Arguments:\n      schedule: a function that takes an epoch index\n          (integer, indexed from 0) and current learning rate\n          as inputs and returns a new learning rate as output (float).\n  \"\"\"\n\n    def __init__(self, schedule):\n        super(CustomLearningRateScheduler, self).__init__()\n        self.schedule = schedule\n\n    def on_epoch_begin(self, epoch, logs=None):\n        if not hasattr(self.model.d_opt, \"lr\"):\n            raise ValueError('Disc_optimizer must have a \"lr\" attribute.')\n        if not hasattr(self.model.g_opt, \"lr\"):\n            raise ValueError('Gen_optimizer must have a \"lr\" attribute.')\n        # Get the current learning rate from model's optimizer.\n        d_lr = float(tf.keras.backend.get_value(self.model.d_opt.learning_rate))\n        g_lr = float(tf.keras.backend.get_value(self.model.g_opt.learning_rate))\n        # Call schedule function to get the scheduled learning rate.\n        res = self.schedule(epoch, d_lr, g_lr)\n        d_scheduled_lr = res[0]\n        g_scheduled_lr = res[1]\n        # Set the value back to the optimizer before this epoch starts\n        tf.keras.backend.set_value(self.model.d_opt.lr, d_scheduled_lr)\n        tf.keras.backend.set_value(self.model.g_opt.lr, g_scheduled_lr)\n        print(\"\\nEpoch %05d: Learning rate is %6.4f (Disc) and %6.4f (Gen).\" % (epoch, d_scheduled_lr, g_scheduled_lr))\n\n\nLR_SCHEDULE = [\n    # (epoch to start, gen learning rate, disc learning rate) tuples\n    # start: (0, 0.01, 0.005) G,D\n    (10, 0.005, 0.0025),\n    (20, 0.0025, 0.00125),\n    (30, 0.00125, 0.000625),\n    (40, 0.000625, 0.0003125)\n]\n\n\ndef lr_schedule(epoch, d_lr, g_lr):\n    \"\"\"Helper function to retrieve the scheduled learning rate based on epoch.\"\"\"\n    if epoch < LR_SCHEDULE[0][0] or epoch > LR_SCHEDULE[-1][0]:\n        return [d_lr, g_lr]\n    for i in range(len(LR_SCHEDULE)):\n        if epoch == LR_SCHEDULE[i][0]:\n            return LR_SCHEDULE[i][1:3]\n    return [d_lr, g_lr]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Training loop**"},{"metadata":{"trusted":true},"cell_type":"code","source":"history = History()\nmonet_gan.fit(\n    monet_ds,\n    #validation_split=0.20,\n    epochs=EPOCHS,\n    callbacks=[\n        checkpoint,\n        GANMonitor(num_img=3, latent_dim=LATENT_DIM),\n        CustomLearningRateScheduler(lr_schedule),\n        history\n    ]\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# summarize losses\nhistory.history['g_loss']\nhistory.history['d_loss']\nplt.plot(history.history['g_loss'])\nplt.plot(history.history['d_loss'])\nplt.title('model loss')\nplt.ylabel('loss')\nplt.xlabel('epoch')\nplt.legend(['generator', 'discriminator'], loc='upper left')\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Display image from epoch-number**"},{"metadata":{"trusted":true},"cell_type":"code","source":"def display_image(num_img, epoch_no):\n  return PIL.Image.open('generated/generated_img_{i}_{epoch}.png'.format(i = num_img, epoch = epoch_no))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# epoch_no from 0 to 49\n# choose num_img from 0 to 2. if want more samples, in monet_gan.fit(), set GAN_monitor(num_img=n,) for your choice of n\ndisplay_image(0, 0) ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Create gif**"},{"metadata":{"trusted":true},"cell_type":"code","source":"anim_file = 'dcgan.gif'\n\nwith imageio.get_writer(anim_file, mode='I') as writer:\n    filenames = glob.glob('generated/generated_img_2*.png') # can also try img_1 or img_2\n    filenames = sorted(filenames)\n    for filename in filenames:\n        image = imageio.imread(filename)\n        writer.append_data(image)\n    image = imageio.imread(filename)\n    writer.append_data(image)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install git+https://github.com/tensorflow/docs\nimport tensorflow_docs.vis.embed as embed\nembed.embed_file(anim_file)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Predict images"},{"metadata":{"trusted":true},"cell_type":"code","source":"from scipy.linalg import sqrtm\ndef calculate_fid(model, images1, images2):\n    # calculate activations\n    #act1 = model.predict(images1)\n    #act2 = model.predict(images2)\n    # calculate mean and covariance statistics\n    mu1, sigma1 = images1.mean(axis=0), np.cov(images1, rowvar=False)\n    mu2, sigma2 = images2.mean(axis=0), np.cov(images2, rowvar=False)\n    # calculate sum squared difference between means\n    ssdiff = np.sum((mu1 - mu2)**2.0)\n    # calculate sqrt of product between cov\n    covmean = sqrtm(sigma1.dot(sigma2))\n    # check and correct imaginary numbers from sqrt\n    if np.iscomplexobj(covmean):\n        covmean = covmean.real\n    # calculate score\n    fid = ssdiff + np.trace(sigma1 + sigma2 - 2.0 * covmean)\n    return fid","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image_dir =  '/images'\nif not os.path.exists(image_dir):\n    os.makedirs(image_dir)\n\nfor i in range(10):\n    prediction = monet_generator(np.random.randn(1, LATENT_DIM), training=False)[0].numpy()\n    prediction = (prediction * 127.5 + 127.5).astype(np.uint8)\n    im = PIL.Image.fromarray(prediction)\n    #print(calculate_fid(monet_gan, np.asarray(im.convert('L')).astype(np.uint32), np.asarray(im.convert('L')).astype(np.uint32)))\n    im.save(f\"/images/{i}.jpg\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fid = 999999999\nprediction = monet_generator(np.random.randn(1, LATENT_DIM), training=False)[0].numpy()\nprediction = (prediction * 127.5 + 127.5).astype(np.uint8)\nprediction = PIL.Image.fromarray(prediction)\nbatch = next(iter(monet_ds))\nfor image in batch:\n    image = image.numpy()\n    image = np.dot(image, [0.2989, 0.5870, 0.1140])\n    t = calculate_fid(monet_gan, np.asarray(prediction.convert('L')).astype(np.uint32), image.astype(np.uint32))\n    if t < fid:\n        fid = t\nprint(fid)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"shutil.make_archive(\"/kaggle/working/images\", 'zip', \"/kaggle/working\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Todo - implement MiFID metric for evaluation, and maybe others (IS?)<br>\nDebug - Spectral Normalization"}],"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}