{"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":"Becoming an artist takes a lot of skill. And becoming an artist with the status enjoyed by [Claude Monet](https://en.wikipedia.org/wiki/Claude_Monet) takes an _incredible_ amount of skill. Hundreds of artists practice their craft for years and years in order to produce something slightly close to Monet.\n\nLuckily, we have Deep Learning. 😉\n\nWhat artists take years to accomplish, we can accomplish in a few hours (even less if we have the right hardware) thanks to the awesome power granted to us by deep learning in the form of [Neural Style Transfer](https://en.wikipedia.org/wiki/Neural_style_transfer) (NST).\n\nNST is a problem in computer vision where the goal is to take an input image and transfer the \"style\" of another image to this image. That is, we move the style of the input image from its original (input) domain to a new (output) domain. The two domains can be anything. In our case, our input domain consists of real-life photos taken by a camera and our output domain consists of Monet paintings. The goal is to make these real-life photos look like Monet paintings.\n\nWhile the problem of NST has been studied for decades, deep learning models made their way into the field with the introduction of the [NST model](https://arxiv.org/abs/1508.06576) in 2015. The model treats the problem as a supervised learning task, i.e. it needs image pairs as input with one image being in the input domain and other being its output-domain counterpart so that it has \"input features\" and the ground truth \"label\". This involves a lot of effort in data collection since for each training image, we also need the expected output image.\n\nLuckily, we have Deep Learning. 😉\n\nThe introduction of [Generative Adversarial Networks](https://en.wikipedia.org/wiki/Generative_adversarial_network) (GANs) in 2014 and the subsequent development of [CycleGAN](https://arxiv.org/abs/1703.10593), introduced in 2017, has allowed the treatment of NST as an unsupervised task, thereby eliminating the need of having paired data.\n\nIn this notebook, we will train a CycleGAN model from scratch on a TPU using [TensorFlow](https://github.com/tensorflow/tensorflow) to transform real-life photos to Monet-esque paintings.\n\n> Note: For a quick overview on how GANs work, see this video: [A Friendly Introduction to Generative Adversarial Networks (GANs)](youtube.com/watch?v=8L11aMN5KY8).\n\n> Note: For a nice explanation of the main ideas in the CycleGAN paper, see this video: [CycleGAN Paper Walkthrough](https://www.youtube.com/watch?v=5jziBapziYE).","metadata":{}},{"cell_type":"markdown","source":"# CycleGAN - Main Ideas","metadata":{}},{"cell_type":"markdown","source":"CycleGAN is based on two key ideas:\n- IST should be \"cycle consistent.\" That is, if we convert an image in domain A to one in domain B and then convert it back into domain A, we should get back the original image. That is, the generators should be able to reverse each others operations.\n- If an image in domain A is given as input to the generator which generates images in domain A, the generator should do nothing since the image already in the generator's domain. That is, the generator should act as an identity function, yielding the image back as the output. The same goes for the generator for domain B.\n\nBoth these ideas lead to a cyclic forward pass. This is why the model is called a _CycleGAN_.\n\nTo achieve this, there are two sets of GANs. That is, we have two generators and two discriminators:\n- There is a generator that takes images from domain A and generates images in domain B\n- There is a discriminator which classifies images in domain B as real or fake\n- There is a generator that takes images from domain B and generates images in domain A\n- There is a discriminator which classifies images in domain A as real or fake\n\nThus, one GAN is responsible for converting an image in domain A to one in domain B and another GAN is responsible for converting an image in domain B to one in domain A. Both these GANs are trained simultaneously and the performance of one contributes in that of the other.\n\nAll discriminators are trained to output the probability of an image being real. An output of 1 implies that the discriminator thinks the image is a real image (not generated) and an output of 0 implies that it thinks the image is fake (generated).\n\nFor this dataset, the following diagram summarises the forward pass through the CycleGAN. The text in `()` denotes the domain each output/input is in (follow the color coded arrows and see \"CycleGAN Model\" section):","metadata":{"_kg_hide-input":true,"_kg_hide-output":true}},{"cell_type":"markdown","source":"![CycleGAN Forward Pass](https://drive.google.com/uc?export=view&id=1EvPhRtsgOoD5WSy8ETsg91-Zl8YCNbkN)","metadata":{"_kg_hide-input":true,"_kg_hide-output":true}},{"cell_type":"markdown","source":"To quantify the performance of the generators and discriminators, the following loss functions are involved:\n\n- **Discriminator Loss** - This takes the outputs of a discriminator for a generated image and the real image, and applies a binary cross entropy loss on them since the discriminator is a binary classifier. The `y_true` is all 0's and all 1's respectively since the first is a fake image and the second is a real image.\n- **Generator Loss** - This takes the output of the discriminator for its generated image and applies a binary cross entropy on it, with the `y_true` being all 1's since the generator's goal is to fool the discriminator into thinking that the generated image is real.\n- **Cycle Consistency Loss* (CCL)* - This takes the original image and the cycled image, and calculates how far the cycled image is from the original image in order to measure the cycle consistency property mentioned above.\n- **Identity Loss** - This takes the original image and the identity image, and calculates how far the identity image is from the original image in order to measure the identity function propery mentioned above.\n\nThe total loss for a generator is the sum of the generator loss, the _total_ (across both the generators) CCL and the identity loss. The total CCL is taken for each generator so that the overall model is cycle consistent as a whole rather one GAN individually.\n\nThe total loss for a discriminator is simply the discriminator loss.\n\nAll in all, we get a total of four loss values after every forward step:\n- Domain A Generator Loss - Loss for the generator which generates images in domain A.\n- Domain B Generator Loss - Loss for the generator which generates images in domain B.\n- Domain A Discriminator Loss - Loss for the discriminator which classifies images in domain A as fake or real.\n- Domain B Discriminator Loss - Loss for the discriminator which classifies images in domain B as fake or real.\n\nThe following diagram summarises the loss calculation for this dataset:","metadata":{"_kg_hide-input":true,"_kg_hide-output":true}},{"cell_type":"markdown","source":"![CycleGAN Loss Calculation](https://drive.google.com/uc?export=view&id=1w5YHLl_XXUOfhMW6nMutbfUX4AwCO6gL)","metadata":{"_kg_hide-input":true,"_kg_hide-output":true}},{"cell_type":"markdown","source":"# Imports","metadata":{}},{"cell_type":"code","source":"# Turn off annoying TensorFlow warnings\nimport os\nos.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"4\"","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-06-21T09:25:49.872723Z","iopub.execute_input":"2022-06-21T09:25:49.873279Z","iopub.status.idle":"2022-06-21T09:25:49.877537Z","shell.execute_reply.started":"2022-06-21T09:25:49.873246Z","shell.execute_reply":"2022-06-21T09:25:49.876611Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import io\nimport os\nimport random\nimport warnings\nimport zipfile\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport tensorflow as tf\nimport tensorflow_addons as tfa\nfrom kaggle_datasets import KaggleDatasets\nfrom PIL import Image\n\nfrom tensorflow.keras import (\n    Input,\n    Model,\n    layers,\n    optimizers,\n    losses,\n    utils\n)\n\nwarnings.filterwarnings(\"ignore\")\n\n%matplotlib inline","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-06-21T11:12:59.711990Z","iopub.execute_input":"2022-06-21T11:12:59.712525Z","iopub.status.idle":"2022-06-21T11:12:59.724321Z","shell.execute_reply.started":"2022-06-21T11:12:59.712473Z","shell.execute_reply":"2022-06-21T11:12:59.723249Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def seed_everything(seed=42):\n    tf.random.set_seed(42)\n    np.random.seed(seed)\n    random.seed(seed)\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n    \nseed_everything()","metadata":{"execution":{"iopub.status.busy":"2022-06-21T09:26:13.237090Z","iopub.execute_input":"2022-06-21T09:26:13.238226Z","iopub.status.idle":"2022-06-21T09:26:13.245793Z","shell.execute_reply.started":"2022-06-21T09:26:13.238165Z","shell.execute_reply":"2022-06-21T09:26:13.244536Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Detect TPU","metadata":{}},{"cell_type":"code","source":"try:\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver.connect()\n    strategy = tf.distribute.TPUStrategy(tpu)\n    print(\"Using TPU:\")\nexcept ValueError:\n    print(\"No TPU found. Falling back to GPU/CPU:\")\n    strategy = tf.distribute.MirroredStrategy()\n\nprint(\"Number of accelerators: \", strategy.num_replicas_in_sync)","metadata":{"execution":{"iopub.status.busy":"2022-06-21T09:26:16.447579Z","iopub.execute_input":"2022-06-21T09:26:16.448447Z","iopub.status.idle":"2022-06-21T09:26:16.496211Z","shell.execute_reply.started":"2022-06-21T09:26:16.448409Z","shell.execute_reply":"2022-06-21T09:26:16.495510Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Configuration\n\nWe will define some basic configuration that will be used throughout the notebook.","metadata":{}},{"cell_type":"code","source":"# This gives us the URL on Google Cloud Storage where this dataset is stored\nGCS_DS_PATH = KaggleDatasets().get_gcs_path(\"gan-getting-started\")\nGCS_DS_PATH","metadata":{"execution":{"iopub.status.busy":"2022-06-21T09:26:47.996645Z","iopub.execute_input":"2022-06-21T09:26:47.997095Z","iopub.status.idle":"2022-06-21T09:26:48.625262Z","shell.execute_reply.started":"2022-06-21T09:26:47.997062Z","shell.execute_reply":"2022-06-21T09:26:48.624061Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Config:\n    DATA_DIR = GCS_DS_PATH\n    \n    AUTOTUNE = tf.data.experimental.AUTOTUNE\n    \n    IMG_SIZE = 256\n    \n    PREFIXES = {\n        \"monet\": \"monet_\",\n        \"photo\": \"photo_\"\n    }\n    \n    # Paper uses a batch size of 1\n    BATCH_SIZE = 1\n    \n    EPOCHS = 50\n\n    LR = 2e-4\n    \n    # This is a scaling factor for the cycle consistency\n    # And the identity losses\n    LAMBDA = 10\n    \n    VAL_SPLIT = 0.1\n    \n    @classmethod\n    def filepath(cls, prefix=\"photo\"):\n        prefixes = tuple(cls.PREFIXES)\n        \n        if prefix not in prefixes:\n            raise ValueError(f\"Unrecognized prefix. Should be in {prefixes}.\")\n            \n        return os.path.join(cls.DATA_DIR, f\"{cls.PREFIXES[prefix]}tfrec\")","metadata":{"execution":{"iopub.status.busy":"2022-06-21T09:26:50.951385Z","iopub.execute_input":"2022-06-21T09:26:50.951856Z","iopub.status.idle":"2022-06-21T09:26:50.960148Z","shell.execute_reply.started":"2022-06-21T09:26:50.951820Z","shell.execute_reply":"2022-06-21T09:26:50.959160Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset Functions","metadata":{}},{"cell_type":"code","source":"# Decode a single image in a TFRecord\n# It also standardizes the image\ndef decode_image(image, channels=3) -> tf.Tensor:\n    img = tf.image.decode_jpeg(image, channels=channels)\n    img = (tf.cast(img, tf.float32) / 127.5) - 1\n    img = tf.reshape(img, [Config.IMG_SIZE, Config.IMG_SIZE, channels])\n    return img","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-06-21T09:26:52.901504Z","iopub.execute_input":"2022-06-21T09:26:52.902792Z","iopub.status.idle":"2022-06-21T09:26:52.908714Z","shell.execute_reply.started":"2022-06-21T09:26:52.902726Z","shell.execute_reply":"2022-06-21T09:26:52.907663Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Read one TFRecord\n# There are no labels and so we only return the image\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","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-06-21T09:27:48.248453Z","iopub.execute_input":"2022-06-21T09:27:48.248900Z","iopub.status.idle":"2022-06-21T09:27:48.256224Z","shell.execute_reply.started":"2022-06-21T09:27:48.248865Z","shell.execute_reply":"2022-06-21T09:27:48.254735Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Create TensorFlow dataset from the given list of filepaths\ndef get_dataset(filepaths, ordered=False) -> tf.data.Dataset:\n    options = tf.data.Options()\n    if ordered is False:\n        options.experimental_deterministic = False\n\n    dataset = tf.data.TFRecordDataset(filepaths, num_parallel_reads=Config.AUTOTUNE)\n    dataset = dataset.with_options(options)\n\n    dataset = dataset.map(read_tfrecord, num_parallel_calls=Config.AUTOTUNE)\n    \n    # Cache helps speed up data access by preventing multipe file I/O operations\n    dataset = dataset.cache()\n    \n    dataset = dataset.shuffle(2048)\n    \n    return dataset","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-06-21T09:27:50.646152Z","iopub.execute_input":"2022-06-21T09:27:50.646667Z","iopub.status.idle":"2022-06-21T09:27:50.654982Z","shell.execute_reply.started":"2022-06-21T09:27:50.646627Z","shell.execute_reply":"2022-06-21T09:27:50.654131Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Example","metadata":{}},{"cell_type":"code","source":"# Load dataset\npath = os.path.join(Config.filepath(\"monet\"), \"*.tfrec\")\nmonet_filenames = tf.io.gfile.glob(path)\n\npath = os.path.join(Config.filepath(\"photo\"), \"*.tfrec\")\nphoto_filenames = tf.io.gfile.glob(path)\n\nmonet_ds = get_dataset(monet_filenames).batch(1)\nphoto_ds = get_dataset(photo_filenames).batch(1)","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-06-21T11:13:09.140741Z","iopub.execute_input":"2022-06-21T11:13:09.141224Z","iopub.status.idle":"2022-06-21T11:13:09.496889Z","shell.execute_reply.started":"2022-06-21T11:13:09.141184Z","shell.execute_reply":"2022-06-21T11:13:09.496004Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Create iterators\nexample_monet = next(iter(monet_ds))\nexample_photo = next(iter(photo_ds))","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-06-21T11:13:10.553103Z","iopub.execute_input":"2022-06-21T11:13:10.553603Z","iopub.status.idle":"2022-06-21T11:13:15.104915Z","shell.execute_reply.started":"2022-06-21T11:13:10.553555Z","shell.execute_reply":"2022-06-21T11:13:15.103857Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Shape: Each image is an RGB image with dimensions 256x256\nexample_monet.shape, example_photo.shape","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-06-21T11:13:15.107060Z","iopub.execute_input":"2022-06-21T11:13:15.107458Z","iopub.status.idle":"2022-06-21T11:13:15.114486Z","shell.execute_reply.started":"2022-06-21T11:13:15.107423Z","shell.execute_reply":"2022-06-21T11:13:15.113141Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Plot a Monet painting and a photo\nmonet = example_monet[0]\nimg = example_photo[0]\n\nf, axs = plt.subplots(1, 2, figsize=(7, 7))\n\naxs[0].imshow(monet)\naxs[0].set_axis_off()\naxs[0].set_title(\"Monet\")\n\naxs[1].imshow(img)\naxs[1].set_axis_off()\naxs[1].set_title(\"Photo\")","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-06-21T11:13:15.115946Z","iopub.execute_input":"2022-06-21T11:13:15.116265Z","iopub.status.idle":"2022-06-21T11:13:15.378242Z","shell.execute_reply.started":"2022-06-21T11:13:15.116235Z","shell.execute_reply":"2022-06-21T11:13:15.377191Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Discriminator\n\nThe discrimnators in a CycleGAN have the following architecture:","metadata":{}},{"cell_type":"markdown","source":"![CycleGan Discriminator](https://drive.google.com/uc?export=view&id=1RJeYlKABS8w9wuKDAzCx9Dyzb4x3opwF)","metadata":{"_kg_hide-input":true,"_kg_hide-output":true}},{"cell_type":"markdown","source":"We will first define a convolutional block (yellow and pink block in the architecture above), which is made up of the following components:\n- `Conv2D` layer\n- Optional `InstanceNormalization`\n- `LeakyReLU` activation\n\nIt takes the input for the layer, the number of filters for `Conv2D`, the stride and a boolean indicating whether instance normalization should be added or not.\n\nThe kernel size is set to 4 and there is a zero padding before the `Conv2D` layer.\n\n> Note: The paper uses mirror or reflection padding but TensorFlow TPU doesn't support that operation. Therefore, reflection padding have been replaced by zero padding.","metadata":{}},{"cell_type":"code","source":"def discriminator_convblock(ip, filters, strides=1, *, norm=True):\n    x = layers.ZeroPadding2D(padding=1)(ip)\n    \n    x = layers.Conv2D(filters=filters, kernel_size=4, strides=strides)(x)\n    \n    if norm is True:\n        x = tfa.layers.InstanceNormalization()(x)\n        \n    x = layers.LeakyReLU()(x)\n    \n    return x","metadata":{"execution":{"iopub.status.busy":"2022-06-21T09:28:52.151062Z","iopub.execute_input":"2022-06-21T09:28:52.151501Z","iopub.status.idle":"2022-06-21T09:28:52.159265Z","shell.execute_reply.started":"2022-06-21T09:28:52.151470Z","shell.execute_reply":"2022-06-21T09:28:52.158395Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Now, we will define the discriminator itself. It takes the number of input channels in the image (defaulting to 3 for RGB images) and an optional list of integers denoting the number of filters for the convolutional blocks before the output layer. If not provided, it uses the values in the architecture above (4 blocks of sizes 64, 128, 256 and 512). Note that the first hidden layer does not have instance normalization.\n\nThe output layer is a `Conv2D` layer with 1 output channel since the discriminator does binary classification (1 - Real, 0 - Fake). Sigmoid is not applied here since it will be applied as part of the loss function.\n\nFor an $N\\times3\\times256\\times256$ image, the output is always $N\\times1\\times30\\times30$, where $N$ is the batch size. Notice that the architecture refers to the discriminator as a \"PatchGAN\". This is beacuse each pixel in this $30\\times30$ output corresponds to a $70\\times70$ patch in the input, i.e. each pixel, after applying sigmoid, is the probability of the corresponding $70\\times70$ patch in the input being real.\n\n> Note: Since the entire network is made up of convolutional layers and has no component that depends on input size, it can handle any input size.","metadata":{}},{"cell_type":"code","source":"def Discriminator(channels=3, conv_filters=None):\n    ip = Input(shape=(Config.IMG_SIZE, Config.IMG_SIZE, channels))\n    \n    # If not provided, set conv_filers to default value\n    if conv_filters is None:\n        conv_filters = [64, 128, 256, 512]\n        \n    filters = conv_filters[-1]\n    \n    x = discriminator_convblock(ip=ip, filters=filters, strides=2, norm=False)\n    \n    for filters in conv_filters[1:-1]:\n        x = discriminator_convblock(ip=x, filters=filters, strides=2)\n        \n    x = discriminator_convblock(ip=x, filters=conv_filters[-1], strides=1)\n    \n    x = layers.ZeroPadding2D(padding=1)(x)\n    \n    op = layers.Conv2D(filters=1, kernel_size=4, strides=1)(x)\n    \n    return Model(inputs=ip, outputs=op)","metadata":{"execution":{"iopub.status.busy":"2022-06-21T09:28:53.882325Z","iopub.execute_input":"2022-06-21T09:28:53.883471Z","iopub.status.idle":"2022-06-21T09:28:53.893432Z","shell.execute_reply.started":"2022-06-21T09:28:53.883425Z","shell.execute_reply":"2022-06-21T09:28:53.892074Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Summary","metadata":{}},{"cell_type":"markdown","source":"We will initialize a Discriminator model and get a nice summary. We see that, with the default values, we have  $3,705,985$ parameters to train.","metadata":{"_kg_hide-input":true,"_kg_hide-output":true}},{"cell_type":"code","source":"d = Discriminator()\nd.summary()","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-06-21T09:28:55.435789Z","iopub.execute_input":"2022-06-21T09:28:55.436403Z","iopub.status.idle":"2022-06-21T09:28:56.188145Z","shell.execute_reply.started":"2022-06-21T09:28:55.436357Z","shell.execute_reply":"2022-06-21T09:28:56.186944Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"utils.plot_model(d)","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-06-21T09:28:58.487248Z","iopub.execute_input":"2022-06-21T09:28:58.487642Z","iopub.status.idle":"2022-06-21T09:28:59.818081Z","shell.execute_reply.started":"2022-06-21T09:28:58.487612Z","shell.execute_reply":"2022-06-21T09:28:59.816551Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Free up RAM\ndel d","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-06-21T09:29:04.767697Z","iopub.execute_input":"2022-06-21T09:29:04.768205Z","iopub.status.idle":"2022-06-21T09:29:04.775764Z","shell.execute_reply.started":"2022-06-21T09:29:04.768158Z","shell.execute_reply":"2022-06-21T09:29:04.774708Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Generator\n\nThe generators in a CycleGAN have the following architecture. It follows a UNet-like architecture with a downsampling part and then an upsampling part, with a little ResNet sprinkled in between.","metadata":{}},{"cell_type":"markdown","source":"![CycleGAN Generator](https://drive.google.com/uc?export=view&id=11bEYU1ZiffmrOeVi5ib7-x8KuSV0KRdb)","metadata":{"_kg_hide-input":true,"_kg_hide-output":true}},{"cell_type":"markdown","source":"We will first define a convolutional block, which can handle both the yellow blocks and the green blocks in the architecture above. It is made up of the following components:\n- Either a `Conv2D` (yellow) layer or a `Conv2DTranspose` (green) layer.\n- Optional `InstanceNormalization`.\n- Optional `ReLU` activation.\n\nIt takes: \n- Input\n- Number of filters for the `Conv2D` or `Conv2DTranspose`\n- Padding\n- Boolean indicating whether to use `Conv2d` or `ConvTranspose2d`\n- Boolean indicating whether to add instance normalization or not\n- Boolean indicating whether to add `ReLU` or not\n- Any additional arguments by keyword. All additional `kwargs` are passed to either `Conv2d` or `ConvTranspose2d` to set things like the kernel size, stride and padding. Not passing `kwargs` will set these to their PyTorch defaults.","metadata":{}},{"cell_type":"code","source":"def generator_convblock(\n    ip,\n    filters,\n    padding=1,\n    *,\n    transpose=False,\n    use_activation=True,\n    norm=True,\n    **kwargs,\n):\n    x = layers.ZeroPadding2D(padding=padding)(ip)\n        \n    klass = layers.Conv2DTranspose if transpose is True else layers.Conv2D\n    \n    if transpose is True:\n        kwargs[\"padding\"] = \"same\"\n    \n    x = klass(filters=filters, **kwargs)(x)\n    \n    if norm is True:\n        x = tfa.layers.InstanceNormalization()(x)\n        \n    if use_activation is True:\n        x = layers.ReLU()(x)\n        \n    return x","metadata":{"execution":{"iopub.status.busy":"2022-06-21T09:29:11.974496Z","iopub.execute_input":"2022-06-21T09:29:11.975583Z","iopub.status.idle":"2022-06-21T09:29:11.984260Z","shell.execute_reply.started":"2022-06-21T09:29:11.975519Z","shell.execute_reply":"2022-06-21T09:29:11.983445Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Next, we will define the residual blocks (orange blocks in the architecture above). These are made up of two of the above convolutional blocks with `Conv2D` as their convolution operation. The first one has an activation while the second one does not. They use a kernel size of 3, stride of 1, padding of 1. They have the same number of input and output channels (since they are residual blocks and use skip connections which involve element-wise addition).","metadata":{}},{"cell_type":"code","source":"def residual_block(ip, filters):\n    x = generator_convblock(\n        ip=ip,\n        filters=filters,\n        kernel_size=3,\n        strides=1,\n        padding=1,\n    )\n    \n    x = generator_convblock(\n        ip=x,\n        filters=filters,\n        kernel_size=3,\n        strides=1,\n        padding=1,\n        use_activation=False,\n    )\n\n    # Skip connection\n    x = layers.Add()([ip, x])\n    \n    return x","metadata":{"execution":{"iopub.status.busy":"2022-06-21T09:29:19.249856Z","iopub.execute_input":"2022-06-21T09:29:19.250618Z","iopub.status.idle":"2022-06-21T09:29:19.256924Z","shell.execute_reply.started":"2022-06-21T09:29:19.250576Z","shell.execute_reply":"2022-06-21T09:29:19.255722Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Finally, we will define the generator itself. It takes the following arguments:\n- Number of input channels in the image (defaulting to 3 for RGB images).\n- Optional list of integers containing the number of filters for the convolutional blocks with `Conv2D` as their convolutional operation. If not provided, it uses the values from the paper (3 blocks of size 64, 128 and 256).\n- Optional filters for the residual blocks (defaulting to 256). This and the last value in the above argument should match .\n- Optional list of integers containing the number of filter for the convolutional blocks with `Conv2DTranspose` as their convolutional operation. If not provided, it uses the values from the paper (2 blocks of size 128 and 64).\n- Number of residual blocks (defaulting to 6).\n\nNote that the first convolutional block has no instance normalization.\n\nThe output layer is a `Conv2D` layer with the same number of filters as the input channels in the image since the generator is responsible for generating another image. It has a `tanh` activation, which makes the output similar to a standardized image.\n\nFor an $N\\times3\\times H\\times W$ image, the output is also $N\\times3\\times H\\times W$.\n\n> Note: A mistake I made earlier was add an `InstanceNormalization` layer after the output layer. This completely crippled the model and it was not learning at all. Do not make the same mistake. 🙂","metadata":{}},{"cell_type":"code","source":"def Generator(\n    channels=3,\n    conv_filters=None,\n    res_filters=256,\n    transpose_filters=None,\n    n_residuals=6,\n):    \n    # If not provided, set conv_filters to the default value\n    if conv_filters is None:\n        conv_filters = [64, 128, res_filters]\n    # If provided, make sure that the last value is same as res_filters\n    elif conv_filters[-1] != res_filters:\n        msg = (\n            f\"Make sure that the last value (={conv_filters[-1]}) \"\n            \"in conv_out_channels is the same as \" \n            f\"res_out_channels (={res_filters})\"\n        )\n        raise ValueError(msg)\n\n    # If not provided, set transpose_filters to default value\n    if transpose_filters is None:\n        transpose_filters = [128, 64]\n        \n    ip = Input(shape=(Config.IMG_SIZE, Config.IMG_SIZE, channels))\n    \n    filters = conv_filters[0]\n    \n    x = generator_convblock(\n        ip=ip,\n        filters=filters,\n        kernel_size=7,\n        strides=1,\n        padding=3,\n        norm=False\n    )\n    \n    for filters in conv_filters[1:]:\n        x = generator_convblock(\n            ip=x,\n            filters=filters,\n            kernel_size=3,\n            strides=2,\n            padding=1,\n        )\n        \n    for _ in range(n_residuals):\n        x = residual_block(ip=x, filters=res_filters)\n        \n    for filters in transpose_filters:\n        x = generator_convblock(\n            ip=x,\n            filters=filters,\n            kernel_size=3,\n            strides=2,\n            padding=0,\n            output_padding=1,\n            transpose=True,\n        )\n        \n        \n    x = layers.ZeroPadding2D(padding=3)(x)\n    # A tanh activation is added\n    # This keeps the input in the range [-1, 1]\n    # Which is similar to standardizing the image\n    op = layers.Conv2D(\n        filters=channels,\n        kernel_size=7,\n        strides=1,\n        activation=\"tanh\",\n    )(x)    \n    \n    return Model(inputs=ip, outputs=op)","metadata":{"execution":{"iopub.status.busy":"2022-06-21T09:29:20.987827Z","iopub.execute_input":"2022-06-21T09:29:20.988210Z","iopub.status.idle":"2022-06-21T09:29:21.002345Z","shell.execute_reply.started":"2022-06-21T09:29:20.988180Z","shell.execute_reply":"2022-06-21T09:29:21.001400Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Summary","metadata":{}},{"cell_type":"markdown","source":"We will initialize a Generator model and get a nice summary. We see that, with the default values, we have  $7,844,995$ parameters to train.","metadata":{"_kg_hide-input":true,"_kg_hide-output":true}},{"cell_type":"code","source":"g = Generator()\ng.summary()","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-06-21T09:29:22.782493Z","iopub.execute_input":"2022-06-21T09:29:22.783108Z","iopub.status.idle":"2022-06-21T09:29:23.546267Z","shell.execute_reply.started":"2022-06-21T09:29:22.783074Z","shell.execute_reply":"2022-06-21T09:29:23.544771Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"utils.plot_model(g)","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-06-21T09:29:24.928835Z","iopub.execute_input":"2022-06-21T09:29:24.929280Z","iopub.status.idle":"2022-06-21T09:29:25.476593Z","shell.execute_reply.started":"2022-06-21T09:29:24.929248Z","shell.execute_reply":"2022-06-21T09:29:25.475282Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Free up RAM\ndel g","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-06-21T09:29:26.032832Z","iopub.execute_input":"2022-06-21T09:29:26.033402Z","iopub.status.idle":"2022-06-21T09:29:26.040192Z","shell.execute_reply.started":"2022-06-21T09:29:26.033352Z","shell.execute_reply":"2022-06-21T09:29:26.039093Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Loss Functions\n\nWe will define a function to calculate the loss for a discriminator.\n\nThe loss for a discriminator is the sum of the binary cross entropy loss for the real image (image which has not been generated) and the fake image (image generated using a generator), since the discriminator is a binary classifier.\n\nThe discriminator is set to output the probability of an image being real. A probability of 1 implies that the discriminator thinks the image is real while a probability of 0 implies that it thinks the image is fake. Thus, for the real image, the target is all ones and for a fake image, the target is all zeros.","metadata":{}},{"cell_type":"code","source":"with strategy.scope():\n    def discriminator_loss(real_disc, fake_disc):\n        # Discriminator is a binary classifier\n        # And so uses Binary Cross Entropy\n        loss_fn = losses.BinaryCrossentropy(from_logits=True, reduction=losses.Reduction.NONE)\n\n        # The real image's y_true is all ones\n        real_loss = loss_fn(tf.ones_like(real_disc), real_disc)\n        # The fake image's y_true is all zeros\n        fake_loss = loss_fn(tf.zeros_like(fake_disc), fake_disc)\n\n        return (real_loss + fake_loss) * 0.5","metadata":{"execution":{"iopub.status.busy":"2022-06-21T09:29:59.948805Z","iopub.execute_input":"2022-06-21T09:29:59.949277Z","iopub.status.idle":"2022-06-21T09:29:59.956782Z","shell.execute_reply.started":"2022-06-21T09:29:59.949241Z","shell.execute_reply":"2022-06-21T09:29:59.955809Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We will define the generator loss for a generator. This loss is used by the generator to measure how good it is at fooling the discriminator. The goal of the generator is to make the discriminator output something as close to all 1s as possible for an image it generates. Thus, this loss is measured using binary cross entropy, where the output of the discriminator for the generated image is the prediction and all 1s is the target.","metadata":{}},{"cell_type":"code","source":"with strategy.scope():\n    def generator_loss(fake_disc):\n        loss_fn = losses.BinaryCrossentropy(from_logits=True, reduction=losses.Reduction.NONE)\n        # The goal of the generator is to fool the discriminator into thinking\n        # The generated image is real\n        # Hence, the loss function uses all 1s as y_true\n        return loss_fn(tf.ones_like(fake_disc), fake_disc)","metadata":{"execution":{"iopub.status.busy":"2022-06-21T09:29:51.348999Z","iopub.execute_input":"2022-06-21T09:29:51.349566Z","iopub.status.idle":"2022-06-21T09:29:51.357954Z","shell.execute_reply.started":"2022-06-21T09:29:51.349517Z","shell.execute_reply":"2022-06-21T09:29:51.356527Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Next, we will define the cycle consistency loss (CCL). CCL is the [L1 norm](https://en.wikipedia.org/wiki/Taxicab_geometry) or the mean absolute error (MAE) between the original image and the output after cycling the image back through the generator. This loss tells us how far the cycled image is from the original image. An alternative to this is the [L2 norm](https://en.wikipedia.org/wiki/Euclidean_distance) or mean squared error (MSE) but that has certain issues when there are outliers in the data.\n\n`scale` is a scaling factor for the loss, which can be used to increase or decrease the contribution of the loss in the final generator loss.","metadata":{}},{"cell_type":"code","source":"with strategy.scope():\n    def cycle_consistency_loss(original, cycled, scale):\n        # Use L1 distance as loss function\n        return tf.math.reduce_mean(tf.abs(original - cycled)) * scale","metadata":{"execution":{"iopub.status.busy":"2022-06-21T09:30:02.875642Z","iopub.execute_input":"2022-06-21T09:30:02.876829Z","iopub.status.idle":"2022-06-21T09:30:02.883213Z","shell.execute_reply.started":"2022-06-21T09:30:02.876781Z","shell.execute_reply":"2022-06-21T09:30:02.882353Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Finally, we will define the identity loss. The identity loss is measured in a way similar to CCL. This measures how far away the identity image (image generated by running original image through its own generator) is from the original image. Here, `scale` is halved, which implies that the identity loss will have half as much contribution to the generator loss as CCL.","metadata":{}},{"cell_type":"code","source":"with strategy.scope():\n    def identity_loss(original, identity, scale):\n        # Use L1 distance as loss function\n        return tf.math.reduce_mean(tf.abs(original - identity)) * scale * 0.5","metadata":{"execution":{"iopub.status.busy":"2022-06-21T09:30:05.489566Z","iopub.execute_input":"2022-06-21T09:30:05.490264Z","iopub.status.idle":"2022-06-21T09:30:05.495845Z","shell.execute_reply.started":"2022-06-21T09:30:05.490229Z","shell.execute_reply":"2022-06-21T09:30:05.494862Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# CycleGAN Model\n\nWe will define a CycleGAN model to put everything together. It takes the generator and discriminator for the Monet paintings, and the generator and discriminator for the real-life photos.\n\nThe forward pass takes a real-life photo (`img`) and a Monet painting (`monet_img`), and does the following:\n\n- Step 1\n    - Converts `img` to a Monet-esque painting using the generator for Monet paintings, giving `fake_monet`.\n    - Converts `monet_img` to a real-life photo using the generator for real-life photos, giving `fake_img`.\n- Step 2\n    - Recreates `monet_img` from `fake_img` using the generator for Monet paintings to test how good the generator is at undoing the changes made on the image. This gives `cycled_monet`.\n    - Repeats the same thing, but this time for `img` using the generator for real-life photos and `fake_monet` . This gives `cycled_img`.\n- Step 3\n    - Passes `monet_img` through the generator for Monet paintings to test how good the generator is at recognizing that `monet_img` is from its own domain and does not need any style transfer. This gives `identity_monet`.\n    - Repeats the same thing, but this time for `img` using the generator for real-life photos. This gives `identity_img`.\n- Step 4\n    - Classifies `monet_img` using the discriminator for Monet paintings. The expected output is close to 1 since `monet_img` is not generated. This gives `real_monet_disc`.\n    - Classifies `img` using the discriminator for real-life photos. The expected output is close to 1 since `img` is not generated. This gives `real_img_disc`.\n- Step 5\n    - Classifies `fake_monet` using the discriminator for Monet paintings. The expected output is close to 0 since `fake_monet` is generated. This gives `fake_monet_disc`.\n    - Classifies `fake_img` using the discriminator for real-life photos. The expected output is close to 0 since `fake_img` is generated. This gives `fake_img_disc`.\n\n\nAll the outputs are returned in a dictionary.\n    \n    \nThe cyclic nature of the network can be seen in the above steps.","metadata":{}},{"cell_type":"code","source":"class CycleGAN(Model):\n    def __init__(self, monet_gen, monet_disc, img_gen, img_disc, scale=10):\n        super().__init__()\n        self.monet_gen = monet_gen\n        self.monet_disc = monet_disc\n        self.img_gen = img_gen\n        self.img_disc = img_disc\n        self.scale = scale\n        \n    def compile(\n        self,\n        monet_gen_opt,\n        monet_disc_opt,\n        img_gen_opt,\n        img_disc_opt,\n        gen_loss_fn,\n        disc_loss_fn,\n        id_loss_fn,\n        ccl_fn\n    ):\n        super().compile()\n        \n        self.monet_gen_opt = monet_gen_opt\n        self.monet_disc_opt = monet_disc_opt\n        \n        self.img_gen_opt = img_gen_opt\n        self.img_disc_opt = img_disc_opt\n        \n        self.gen_loss_fn = gen_loss_fn\n        self.disc_loss_fn = disc_loss_fn\n        self.id_loss_fn = id_loss_fn\n        self.ccl_fn = ccl_fn\n        \n    def call(self, data):\n        img, monet_img = data\n        \n        # Generate the fake Monet painting from original image\n        # Generate the fake image from the original Monet\n        fake_monet = self.monet_gen(img)\n        fake_img = self.img_gen(monet_img)\n\n        # Recreate the Monet from the fake image\n        # Recreate the original image from the fake Monet\n        cycled_monet = self.monet_gen(fake_img)\n        cycled_img = self.img_gen(fake_monet)\n\n        # Generate the identity images from the images\n        identity_monet = self.monet_gen(monet_img)\n        identity_img = self.img_gen(img)\n\n        # Use discriminators to classifiy the real images\n        # Expected output is close to 1 since the images are real\n        real_monet_disc = self.monet_disc(monet_img)\n        real_img_disc = self.img_disc(img)\n\n        # Use discriminators to classify the fake images\n        # Expected output is close to 0 since the images are fake\n        fake_monet_disc = self.monet_disc(fake_monet)\n        fake_img_disc = self.img_disc(fake_img)\n        \n        # Store all the outputs related to the normal image\n        img_output = {\n            \"fake_img\": fake_img,\n            \"cycled_img\": cycled_img,\n            \"identity_img\": identity_img,\n            \"real_img_disc\": real_img_disc,\n            \"fake_img_disc\": fake_img_disc,\n        }\n        \n        # Store all the outputs related to the Monet image\n        monet_output = {\n            \"fake_monet\": fake_monet,\n            \"cycled_monet\": cycled_monet,\n            \"identity_monet\": identity_monet,\n            \"real_monet_disc\": real_monet_disc,\n            \"fake_monet_disc\": fake_monet_disc,\n        }\n        \n        # Return all the outputs\n        return {\n            **img_output,\n            **monet_output,\n        }\n        \n    def train_step(self, data):\n        img, monet_img = data\n        \n        with tf.GradientTape(persistent=True) as tape:\n            outputs = self(data)\n            \n            # Calculate the generator losses\n            img_gen_loss = self.gen_loss_fn(outputs[\"fake_img_disc\"])\n            monet_gen_loss = self.gen_loss_fn(outputs[\"fake_monet_disc\"])\n            \n            # Add the identity losses to the generator losses\n            img_gen_loss += self.id_loss_fn(\n                original=img,\n                identity=outputs[\"identity_img\"],\n                scale=self.scale,\n            )\n            monet_gen_loss += self.id_loss_fn(\n                original=monet,\n                identity=outputs[\"identity_monet\"],\n                scale=self.scale,\n            )  \n            \n            # Calculate the total CCL\n            total_ccl = self.ccl_fn(\n                original=img,\n                cycled=outputs[\"cycled_img\"],\n                scale=self.scale,\n            )\n            total_ccl += self.ccl_fn(\n                original=monet,\n                cycled=outputs[\"cycled_monet\"],\n                scale=self.scale,\n            )\n            \n            # Add the total CCL to the generator losses\n            img_gen_loss += total_ccl\n            monet_gen_loss += total_ccl\n            \n            # Calculate the discriminator losses\n            img_disc_loss = self.disc_loss_fn(\n                outputs[\"real_img_disc\"],\n                outputs[\"fake_img_disc\"],\n            )\n            \n            monet_disc_loss = self.disc_loss_fn(\n                outputs[\"real_monet_disc\"],\n                outputs[\"fake_monet_disc\"],\n            )\n            \n        # Calculate gradients for img_gen\n        img_gen_weights = self.img_gen.trainable_weights\n        img_gen_grads = tape.gradient(img_gen_loss, img_gen_weights)\n        \n        # Calculate gradients for monet_gen\n        monet_gen_weights = self.monet_gen.trainable_weights\n        monet_gen_grads = tape.gradient(monet_gen_loss, monet_gen_weights)\n        \n        # Calculate gradients for img_disc\n        img_disc_weights = self.img_disc.trainable_weights\n        img_disc_grads = tape.gradient(img_disc_loss, img_disc_weights)\n        \n        # Calculate gradients for monet_disc\n        monet_disc_weights = self.monet_disc.trainable_weights\n        monet_disc_grads = tape.gradient(monet_disc_loss, monet_disc_weights)\n        \n        # Apply the gradients to the respective optimizers\n        self.img_gen_opt.apply_gradients(zip(img_gen_grads, img_gen_weights))\n        self.monet_gen_opt.apply_gradients(zip(monet_gen_grads, monet_gen_weights))\n        \n        self.img_disc_opt.apply_gradients(zip(img_disc_grads, img_disc_weights))\n        self.monet_disc_opt.apply_gradients(zip(monet_disc_grads, monet_disc_weights))\n        \n        # Return all the loss values\n        return {\n            \"img_gen_loss\": img_gen_loss,\n            \"img_disc_loss\": img_disc_loss,\n            \"monet_gen_loss\": monet_gen_loss,\n            \"monet_disc_loss\": monet_disc_loss,\n        }","metadata":{"execution":{"iopub.status.busy":"2022-06-21T09:30:10.295364Z","iopub.execute_input":"2022-06-21T09:30:10.295806Z","iopub.status.idle":"2022-06-21T09:30:10.322368Z","shell.execute_reply.started":"2022-06-21T09:30:10.295772Z","shell.execute_reply":"2022-06-21T09:30:10.321361Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Summary","metadata":{}},{"cell_type":"markdown","source":"We will initialize a `CycleGAN` model and get a nice summary. We see that, with the default values, we have $23,101,960$ parameters to train. This checks out since there are two sets of a generator and a discriminator, and $2*(7,844,995+3,705,985)=23,101,972$","metadata":{"_kg_hide-input":true}},{"cell_type":"code","source":"monet_gen = Generator()\nmonet_disc = Discriminator()\n\nimg_gen = Generator()\nimg_disc = Discriminator()\n\ngan = CycleGAN(monet_gen=monet_gen, monet_disc=monet_disc, img_gen=img_gen, img_disc=img_disc)\ngan.build(input_shape=[(None, 256, 256, 3), (None, 256, 256, 3)])\ngan.summary()","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-06-21T09:30:15.939182Z","iopub.execute_input":"2022-06-21T09:30:15.940047Z","iopub.status.idle":"2022-06-21T09:30:21.212271Z","shell.execute_reply.started":"2022-06-21T09:30:15.940010Z","shell.execute_reply":"2022-06-21T09:30:21.210914Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Free up RAM\ndel monet_gen\ndel monet_disc\ndel img_gen\ndel img_disc\ndel gan","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-06-21T09:31:31.636036Z","iopub.execute_input":"2022-06-21T09:31:31.636640Z","iopub.status.idle":"2022-06-21T09:31:31.645790Z","shell.execute_reply.started":"2022-06-21T09:31:31.636600Z","shell.execute_reply":"2022-06-21T09:31:31.644626Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training","metadata":{}},{"cell_type":"markdown","source":"## Load Datasets","metadata":{}},{"cell_type":"code","source":"path = os.path.join(Config.filepath(\"monet\"), \"*.tfrec\")\nmonet_filenames = tf.io.gfile.glob(path)\n\npath = os.path.join(Config.filepath(\"photo\"), \"*.tfrec\")\nphoto_filenames = tf.io.gfile.glob(path)","metadata":{"execution":{"iopub.status.busy":"2022-06-21T09:31:34.375772Z","iopub.execute_input":"2022-06-21T09:31:34.376740Z","iopub.status.idle":"2022-06-21T09:31:34.651181Z","shell.execute_reply.started":"2022-06-21T09:31:34.376695Z","shell.execute_reply":"2022-06-21T09:31:34.650119Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"monet_ds = get_dataset(monet_filenames).batch(Config.BATCH_SIZE)\nphoto_ds = get_dataset(photo_filenames).batch(Config.BATCH_SIZE)\n\n# Prefetch allows TensorFlow\n# To fetch data parallely while it is training the model\nmonet_ds = monet_ds.prefetch(Config.AUTOTUNE)\nphoto_ds = photo_ds.prefetch(Config.AUTOTUNE)","metadata":{"execution":{"iopub.status.busy":"2022-06-21T09:31:35.709185Z","iopub.execute_input":"2022-06-21T09:31:35.709713Z","iopub.status.idle":"2022-06-21T09:31:35.769727Z","shell.execute_reply.started":"2022-06-21T09:31:35.709672Z","shell.execute_reply":"2022-06-21T09:31:35.768393Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Initialize CycleGAN\n\nWe will use the default values.","metadata":{}},{"cell_type":"code","source":"with strategy.scope():\n    monet_gen = Generator()\n    monet_disc = Discriminator()\n    \n    img_gen = Generator()\n    img_disc = Discriminator()\n    \n    cycle_gan = CycleGAN(\n        monet_gen=monet_gen,\n        monet_disc=monet_disc,\n        img_gen=img_gen,\n        img_disc=img_disc,\n        scale=Config.LAMBDA,\n    )","metadata":{"execution":{"iopub.status.busy":"2022-06-21T09:31:37.179555Z","iopub.execute_input":"2022-06-21T09:31:37.180017Z","iopub.status.idle":"2022-06-21T09:31:39.254090Z","shell.execute_reply.started":"2022-06-21T09:31:37.179973Z","shell.execute_reply":"2022-06-21T09:31:39.253083Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Initialize Optimizers and Compile the Model","metadata":{}},{"cell_type":"code","source":"with strategy.scope():\n    monet_gen_opt = optimizers.Adam(learning_rate=Config.LR, beta_1=0.5)\n    monet_disc_opt = optimizers.Adam(learning_rate=Config.LR, beta_1=0.5)\n    \n    img_gen_opt = optimizers.Adam(learning_rate=Config.LR, beta_1=0.5)\n    img_disc_opt = optimizers.Adam(learning_rate=Config.LR, beta_1=0.5)\n    \n    cycle_gan.compile(\n        monet_gen_opt=monet_gen_opt,\n        monet_disc_opt=monet_disc_opt,\n        img_gen_opt=img_gen_opt,\n        img_disc_opt=img_disc_opt,\n        gen_loss_fn=generator_loss,\n        disc_loss_fn=discriminator_loss,\n        id_loss_fn=identity_loss,\n        ccl_fn=cycle_consistency_loss,\n    )","metadata":{"execution":{"iopub.status.busy":"2022-06-21T09:31:39.659385Z","iopub.execute_input":"2022-06-21T09:31:39.659998Z","iopub.status.idle":"2022-06-21T09:31:39.693256Z","shell.execute_reply.started":"2022-06-21T09:31:39.659958Z","shell.execute_reply":"2022-06-21T09:31:39.692391Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Get, Set and Go!","metadata":{}},{"cell_type":"code","source":"history = cycle_gan.fit(\n    x=tf.data.Dataset.zip((photo_ds, monet_ds)),\n    epochs=Config.EPOCHS,\n)","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-06-21T09:31:42.244766Z","iopub.execute_input":"2022-06-21T09:31:42.245828Z","iopub.status.idle":"2022-06-21T09:31:57.000046Z","shell.execute_reply.started":"2022-06-21T09:31:42.245790Z","shell.execute_reply":"2022-06-21T09:31:56.997336Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Loss Curve\n\nWe will visualize the loss curves for the generators and discriminators. ","metadata":{}},{"cell_type":"code","source":"# Calculate the mean loss per epoch\nkeys = [\"img_gen_loss\", \"img_disc_loss\", \"monet_gen_loss\", \"monet_disc_loss\"]\n\nepoch_history = {\"img\": {}, \"monet\": {}}\n\nfor key in keys:\n    img_type, model, _ = key.split(\"_\")\n    epoch_history[img_type][model] = np.array(\n        [tf.reduce_mean(loss).numpy() for loss in history.history[key]]\n    )","metadata":{"execution":{"iopub.status.busy":"2022-06-19T21:49:16.249145Z","iopub.execute_input":"2022-06-19T21:49:16.249895Z","iopub.status.idle":"2022-06-19T21:49:16.262406Z","shell.execute_reply.started":"2022-06-19T21:49:16.24985Z","shell.execute_reply":"2022-06-19T21:49:16.261439Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"titles = {\"img\": \"Image\", \"monet\": \"Monet\"}\n\nf, axs = plt.subplots(1, len(epoch_history), figsize=(10, 5))\n\nfor (key, losses), ax in zip(epoch_history.items(), axs.flatten()):\n    ax.plot(epoch_history[key][\"gen\"], label=\"Generator\")\n    ax.plot(epoch_history[key][\"disc\"], label=\"Discriminator\")\n    ax.set_xlabel(\"Epochs\")\n    ax.set_ylabel(\"Loss\")\n    ax.set_title(titles[key])\n    ax.legend()\n\nplt.tight_layout()","metadata":{"execution":{"iopub.status.busy":"2022-06-19T21:49:19.029036Z","iopub.execute_input":"2022-06-19T21:49:19.02936Z","iopub.status.idle":"2022-06-19T21:49:19.39327Z","shell.execute_reply.started":"2022-06-19T21:49:19.029328Z","shell.execute_reply":"2022-06-19T21:49:19.392353Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Visual Validation\n\nWe will plot some of the outputs from the generator for Monet paintings.","metadata":{}},{"cell_type":"code","source":"n_imgs = 5\n\nf, axs = plt.subplots(\n    nrows=2,\n    ncols=n_imgs,\n    figsize=(10, 10),\n    sharex=True,\n    sharey=True,\n)\n\naxs = axs.flatten()\nds = enumerate(photo_ds.take(n_imgs))\n               \n# Plot side-by-side: Photo in even axis and Monet painting in odd axis\nfor (idx, img), ax1, ax2 in zip(ds, axs[::2], axs[1::2]):\n    prediction = cycle_gan.monet_gen(img, training=False)[0].numpy()\n    prediction = (prediction * 127.5 + 127.5).astype(np.uint8)\n    \n    ax1.imshow(img[0])\n    ax1.set_axis_off()\n    ax1.set_title(f\"Photo-{idx + 1}\")\n    \n    ax2.imshow(prediction)\n    ax2.set_axis_off()\n    ax2.set_title(f\"Gen. Monet-{idx + 1}\")\n    \nplt.tight_layout()","metadata":{"execution":{"iopub.status.busy":"2022-06-21T11:11:59.525723Z","iopub.execute_input":"2022-06-21T11:11:59.527723Z","iopub.status.idle":"2022-06-21T11:12:29.588911Z","shell.execute_reply.started":"2022-06-21T11:11:59.527634Z","shell.execute_reply":"2022-06-21T11:12:29.587903Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Submission\n\nWe will create the submission file by directly writing to a ZIP file.","metadata":{}},{"cell_type":"code","source":"from tqdm.notebook import tqdm\n\nwith tqdm(photo_ds, unit=\"imgs\", total=7038) as ds, zipfile.ZipFile(\"images.zip\", \"w\") as zf:\n    for idx, img in enumerate(ds):\n        prediction = cycle_gan.monet_gen(img, training=False)[0].numpy()\n        prediction = (prediction * 127.5 + 127.5).astype(np.uint8)\n        img = Image.fromarray(prediction)\n        dest = io.BytesIO()\n        img.save(dest, format=\"JPEG\")\n        zf.writestr(f\"{idx + 1}.jpg\", dest.getvalue())","metadata":{"execution":{"iopub.status.busy":"2022-06-19T21:55:39.848102Z","iopub.execute_input":"2022-06-19T21:55:39.848438Z","iopub.status.idle":"2022-06-19T21:56:21.972909Z","shell.execute_reply.started":"2022-06-19T21:55:39.848402Z","shell.execute_reply":"2022-06-19T21:56:21.971443Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# References\n\n1. Leon A. Gatys, Alexander S. Ecker, Matthias Bethge. 2015. [A Neural Algorithm of Artistic Style](https://arxiv.org/abs/1508.06576).\n2. Ian J. Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, Yoshua Bengio. 2014. [Generative Adversarial Networks](https://arxiv.org/abs/1406.2661).\n3. Jun-Yan Zhu, Taesung Park, Phillip Isola, Alexei A. Efros. 2017. [Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks](https://arxiv.org/abs/1703.10593).","metadata":{}}]}