{"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":"# How to Paint with Generative Adversarial Networks\n\n*This is part 1, see part 2 [here](https://www.kaggle.com/jesperdramsch/how-to-paint-with-more-gans-ls-cyclegan-aug)*\n\nYou can create Monets with Generative Adversarial Networks (GAN) in a few different ways. We can generate them from scratch using one GAN, where the GAN basically imagines a Monet from scratch. GANs are technically two networks that work against each other, illustrated below. The artist (generator) draws its inspiration from a noise sample and creates a rendering of the data you are trying to generate with said GAN. The private investigator (discriminator) randomly gets assigned real and fake data to investigate. \n\n![GAN from Dramsch PhD thesis.](https://dramsch.net/images/GAN.png)\nFigure describing Generative Adversarial Networks from my [PhD thesis](https://dramsch.net/phd).\n\nThe learning process is collaborative. The generator gets better at fooling the discriminator and the discriminator gets better at figuring out which data is real and which isn't. In mathematical terms they are learning until a [Nash equilibrium](https://en.wikipedia.org/wiki/Nash_equilibrium) is reached, which means neither can learn new tricks and get better. They're a really cool concept and even used in scientific simulation at [CERN](https://indico.cern.ch/event/595059/contributions/2497383/attachments/1431666/2199445/gan_presentation_IML.pdf).\n\nYou can probably guess that they can be tricky to train, due to so many moving parts. This has become a very popular area of research, warranting a [GAN Zoo](https://github.com/hindupuravinash/the-gan-zoo) of all named GANs. Some important stuff you may want to check out if you're interested are keywords like Wasserstein GANs, Gradient Penalization, Attention, and in this context Style Transfer (namely face2face). Maybe you'll even find some in the comments.\n\nIn this tutorial we'll go in depth into:\n\n- Creating images from noise\n- Deep Convolutional Generative Adversarial Networks (DCGANs)\n- Better Loss Functions\n- Data Augmentations\n- Training Tensorflow GANs on TPUs\n\n*This copies in part from my [Intro to Deepfakes](https://www.kaggle.com/jesperdramsch/intro-to-deep-fakes-videos-and-metadata-eda) if you're interested to learn how GANs are used to alter images, videos, and sounds. (Did I mention they're quite versatile?!) This also builds on the Baseline Tutorial. [Please head over and upvote!](https://www.kaggle.com/amyjang/monet-cyclegan-tutorial)*","metadata":{}},{"cell_type":"markdown","source":"## Introduction and Setup\n\nFor this tutorial, we will be using the TFRecord dataset. Import the following packages and change the accelerator to TPU. Because TPUs are pretty awesome.\n\n![](https://i.imgur.com/hRLjugH.png)","metadata":{}},{"cell_type":"code","source":"import tensorflow as tf\nfrom tensorflow import keras\nfrom tensorflow.keras import layers\nimport tensorflow_addons as tfa\nimport tensorflow_datasets as tfds\n\nfrom kaggle_datasets import KaggleDatasets\nimport matplotlib.pyplot as plt\nimport numpy as np\n\nfrom functools import partial\nfrom albumentations import (\n    Compose, RandomBrightness, JpegCompression, HueSaturationValue, RandomContrast, HorizontalFlip,\n    Rotate\n)\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__)","metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Load in the data\n\nWe want to keep our photo dataset and our Monet dataset separate. First, load in the filenames of the TFRecords. We'll load both for the CycleGAN. For the first GAN we only need the Monets as training data.\n\nAll the images for the competition are already sized to `256 x 256`. As these images are RGB images, set the channel to 3. Additionally, we need to scale the images to a `[-1, 1]` scale. Because we are building a generative model, we don't need the labels or the image id so we'll only return the image from the TFRecord.","metadata":{}},{"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))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"You can see I put down a bit of augmentation using `random_jitter` and `flip` to increase our data set, because we simply don't have enough data for","metadata":{}},{"cell_type":"code","source":"IMAGE_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, [256, 256, 3])\n    image = tf.image.decode_jpeg(image, channels=3)\n    #image = (tf.cast(image, tf.float32) / 127.5) - 1\n    image = tf.reshape(image, [*IMAGE_SIZE, 3])\n    return image\n\ndef 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(100000, reshuffle_each_iteration=True).repeat(repeats))\n    dataset = dataset.map(normalize, num_parallel_calls=AUTOTUNE).shuffle(100000)\n    return dataset","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Then load the data and display the first images to see if it all worked out. Which of course it does, because it's taken directly from the tutorial.","metadata":{}},{"cell_type":"code","source":"monet_ds = load_dataset(MONET_FILENAMES, labeled=True, repeats=100).batch(100, drop_remainder=True)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"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)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Build the DCGAN\n## Network Upsample and Downsample\nThe `downsample`, as the name suggests, reduces the 2D dimensions, the width and height, of the image by the stride. The stride is the length of the step the filter takes. Since the stride is 2, the filter is applied to every other pixel, hence reducing the weight and height by 2.\n\nWe'll be using an instance normalization instead of batch normalization. As the instance normalization is not standard in the TensorFlow API, we'll use the layer from TensorFlow Add-ons.","metadata":{}},{"cell_type":"code","source":"def downsample(filters, size, apply_instancenorm=True):\n    initializer = tf.random_normal_initializer(0., 0.02)\n    gamma_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n\n    result = keras.Sequential()\n    result.add(layers.Conv2D(filters, size, padding='same',\n                             kernel_initializer=initializer, use_bias=False))\n    result.add(layers.MaxPool2D())\n\n    if apply_instancenorm:\n        result.add(tfa.layers.InstanceNormalization(gamma_initializer=gamma_init))\n\n    result.add(layers.LeakyReLU())\n\n    return result","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"`Upsample` does the opposite of downsample and increases the dimensions of the of the image. `Conv2DTranspose` does basically the opposite of a `Conv2D` layer.","metadata":{}},{"cell_type":"code","source":"def upsample(filters, size, apply_dropout=False):\n    initializer = tf.random_normal_initializer(0., 0.02)\n    gamma_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n\n    result = keras.Sequential()\n    result.add(layers.Conv2DTranspose(filters, size, strides=2,\n                                      padding='same',\n                                      kernel_initializer=initializer,\n                                      use_bias=False))\n\n    result.add(tfa.layers.InstanceNormalization(gamma_initializer=gamma_init))\n\n    if apply_dropout:\n        result.add(layers.Dropout(0.5))\n\n    result.add(layers.LeakyReLU())\n\n    return result","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Build the generator\n\nThis generator samples from noise, reshapes it and upsamples the entire thing. That's basically it.","metadata":{}},{"cell_type":"code","source":"def Generator(LATENT_DIM=128, OUTPUT_CHANNELS=3):\n    inputs = layers.Input(shape=[LATENT_DIM,])\n\n    up_stack = [\n        upsample(LATENT_DIM, 4, apply_dropout=True), \n        upsample(LATENT_DIM, 4, apply_dropout=True), \n        upsample(LATENT_DIM//2, 4), \n        upsample(LATENT_DIM//4, 4), \n        upsample(LATENT_DIM//8, 4), \n    ]\n\n    initializer = tf.random_normal_initializer(0., 0.02)\n    last = layers.Conv2DTranspose(OUTPUT_CHANNELS, 4,\n                                  strides=2,\n                                  padding='same',\n                                  kernel_initializer=initializer,\n                                  activation='tanh') # (bs, 256, 256, 3)\n    q = 4\n    x = layers.Dense(q * q * LATENT_DIM)(inputs)\n    x = layers.LeakyReLU(alpha=0.2)(x)\n    x = layers.Reshape((q, q, LATENT_DIM))(x)\n\n    # Upsampling \n    for up in up_stack:\n        x = up(x)\n\n    x = last(x)\n\n    return keras.Model(inputs=inputs, outputs=x, name=\"generator\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Build the discriminator\n\nThe discriminator takes in the input image and classifies it as real or fake (generated). Instead of outputing 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.","metadata":{}},{"cell_type":"code","source":"def Discriminator():\n    initializer = tf.random_normal_initializer(0., 0.02)\n    gamma_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n\n    inp = layers.Input(shape=[256, 256, 3], name='input_image')\n\n    x = inp\n\n    down1 = downsample(64, 4, False)(x) # (bs, 128, 128, 64)\n    down2 = downsample(128, 4)(down1) # (bs, 64, 64, 128)\n    down3 = downsample(256, 4)(down2) # (bs, 32, 32, 256)\n\n    zero_pad1 = layers.ZeroPadding2D()(down3) # (bs, 34, 34, 256)\n    conv = layers.Conv2D(512, 4, strides=1,\n                         kernel_initializer=initializer,\n                         use_bias=False)(zero_pad1) # (bs, 31, 31, 512)\n\n    norm1 = tfa.layers.InstanceNormalization(gamma_initializer=gamma_init)(conv)\n\n    leaky_relu = layers.LeakyReLU()(norm1)\n\n    zero_pad2 = layers.ZeroPadding2D()(leaky_relu) # (bs, 33, 33, 512)\n\n    last_conv = layers.Conv2D(1, 4, strides=1,\n                         kernel_initializer=initializer)(zero_pad2) # (bs, 30, 30, 1)\n    last_relu = layers.LeakyReLU(alpha=0.2)(last_conv)\n    last_pool = layers.Flatten()(last_relu)\n    last = layers.Dense(1, activation='sigmoid')(last_pool)\n\n    return tf.keras.Model(inputs=inp, outputs=last)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Define the Least Squares Loss","metadata":{}},{"cell_type":"code","source":"with 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","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Throw it on the TPU\n \nAnd build a GAN with the help of the [keras documentation](https://keras.io/examples/generative/dcgan_overriding_train_step/). Notice that I changed the loss function from the Binary Crossentropy used in the CycleGAN to a Least Squares approach taken from the [RaLSGAN for dogs](https://www.kaggle.com/speedwagon/ralsgan-dogs). The DCGAN would not converge on BCE whatsoever and the LSGAN works surprisingly well for the amount of data we have.","metadata":{}},{"cell_type":"code","source":"OUTPUT_CHANNELS = 3\nLATENT_DIM = 128\n\nwith strategy.scope():\n    monet_generator = Generator(LATENT_DIM, 3) # generates Monet-esque paintings\n    monet_discriminator = Discriminator() # differentiates real Monet paintings and generated Monet paintings","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class 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_real):\n        if isinstance(images_real, tuple):\n            images_real = images_real[0]\n        \n        # Sample random points in the latent space\n        batch_size = tf.shape(images_real)[0]\n        random_latent_vectors = tf.random.normal(shape=(batch_size, self.latent_dim))\n        \n        # labels_gen  = tf.fill((batch_size, 1), self.fake_label)\n        # labels_real = tf.fill((batch_size, 1), self.real_label)\n        \n        labels_gen  = tf.zeros((batch_size, 1)) + self.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 d_tape:\n            # Train the generator\n            images_gen = self.generator(random_latent_vectors, training=False)\n            \n            # Train the discriminator\n            predictions_real = self.discriminator(images_real, training=True)\n            predictions_gen  = self.discriminator(images_gen, training=True)\n                        \n            # Calculate loss LSGan\n            # d_loss = (tf.reduce_mean((predictions_gen  - tf.reduce_mean(predictions_real) + labels_real) ** 2) +\n            #           tf.reduce_mean((predictions_real - tf.reduce_mean(predictions_gen)  - labels_real) ** 2))/2\n            \n            d_loss = self.d_loss_fn(predictions_real, predictions_gen, labels_real)\n            \n        d_grads = d_tape.gradient(d_loss, self.discriminator.trainable_weights)\n        self.d_opt.apply_gradients(zip(d_grads, self.discriminator.trainable_weights))        \n        \n        with tf.GradientTape() as g_tape:\n            # Train the generator\n            images_gen = self.generator(random_latent_vectors, training=True)\n            \n            # Train the discriminator\n            predictions_real = self.discriminator(images_real, training=False)\n            predictions_gen  = self.discriminator(images_gen, training=False)\n                        \n            # Calculate loss LSGan\n            # g_loss = (tf.reduce_mean((predictions_real - tf.reduce_mean(predictions_gen)  + labels_real) ** 2) +\n            #           tf.reduce_mean((predictions_gen  - tf.reduce_mean(predictions_real) - labels_real) ** 2))/2\n            \n            g_loss = self.g_loss_fn(predictions_real, predictions_gen, labels_real)\n        \n        g_grads = g_tape.gradient(g_loss, self.generator.trainable_weights)\n        self.g_opt.apply_gradients(zip(g_grads, self.generator.trainable_weights))\n        \n        return {\"d_loss\": d_loss, \"g_loss\": g_loss}","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Things to note\nWe're using a few tricks here including soft labeling, differing learning rates, augmentation, a better loss function, etc. So this is where we compile and train the model.","metadata":{}},{"cell_type":"code","source":"EPOCHS = 50\n\nLR_G = 0.001\nLR_D = 0.0005\nbeta_1 = .5\n\nreal_label = .66\nfake_label = 0\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    )\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"monet_gan.fit(\n    monet_ds,\n    epochs=EPOCHS,\n)","metadata":{"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"_, ax = plt.subplots(1, 4, figsize=(25, 5))\nfor i in range(4):\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    \n    ax[i].imshow(prediction)\n    ax[i].set_title(\"Monet-esque\")\n    ax[i].axis(\"off\")\nplt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# What about Second GANs?\n\nAlternatively, because us machine learning scientists had a bit too much GPUs running idle, people came up with the idea of cycle-consistent GANs. CycleGANs were introduced for [unpaired Image-to-Image Translation](https://junyanz.github.io/CycleGAN/), for when you don't have Monet available to paint your favorite subject. They're pretty useful generally and have been applied in many domains and style transfer applications. The problem?! This is now two GANs to train that perforn the forward and backward transformation to create the new style from nothing. Finally we can compare apples to oranges.\n\n![](https://junyanz.github.io/CycleGAN/images/objects.jpg)\n\n### [Check that notebook](https://www.kaggle.com/jesperdramsch/how-to-paint-with-more-gans-ls-cyclegan-aug/) out here, where we use the improvements from this notebook!","metadata":{}},{"cell_type":"markdown","source":"# Create submission files\n","metadata":{}},{"cell_type":"code","source":"import PIL\n! mkdir ../images","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for i in range(10000):\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    im.save(f\"../images/{i}.jpg\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import shutil\nshutil.make_archive(\"/kaggle/working/images\", 'zip', \"/kaggle/images\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}