{"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":"<center>\n<div style=\"color:;\n           display:fill;\n           border-radius:5px;\n           background-color: lightgray;\n           font-size:110%;\n           letter-spacing:0.5px\">\n\n<h2 style=\"padding: 10px;\n              color:white;\">Monet: CycleGAN with Better Cycles\n</h2>\n</div>","metadata":{"id":"2DskX8Z8MRfD"}},{"cell_type":"markdown","source":"1. [Introduction](#introduction)\n1. [Data Collection](#datacollection)\n1. [Data Analysis](#dataanalysis)\n    1. [Data Visualization](#datavisualization)\n    1. [EDA Summary](#edasummary)\n1. [Data Preprocessing](#datapreprocessing)\n\n1. [Model Training and Evaluation](#trainandevaluate)\n    1. [Training](#training)\n    1. [Plot predictions](#plotpredictions)\n1. [Submission](#submission)\n1. [Conclusion](#conclusion)\n1. [References](#references)","metadata":{"id":"93y3H1kbMRfK"}},{"cell_type":"markdown","source":"# Introduction <a class=\"anchor\" id=\"introduction\"></a>","metadata":{"id":"x718TA1hMRfL"}},{"cell_type":"markdown","source":"In this notebook I will try to implement approaches from the paper [CycleGAN with Better Cycles](https://www.tongzhouwang.info/better_cycles/report.pdf). This paper proposes some modifications to the CycleGAN framework to improve the quality of the image-to-image translation.\n\nThe main ideas of this paper are to gradually decrease the weight of the cycle consistency loss λ as training progresses and to include an L1 loss on the CNN features extracted by the corresponding discriminator. The modified cycle consistency loss for one direction is defined as a linear combination of CNN feature level and pixel level consistency:\n\n\\begin{aligned}\n& \\tilde{\\mathcal{L}}_{\\mathrm{cyc}}\\left(G, F, D_X, X, \\gamma\\right)= \\mathbb{E}_{x \\sim p_{\\text {data }}(x)}\\left[\\gamma\\left\\|f_{D_X}(F(G(x)))-f_{D_X}(x)\\right\\|_1+(1-\\gamma)\\|F(G(x))-x\\|_1\\right]\n\\end{aligned}\n\nwhere $f_{D_{(\\cdot)}}$ is the feature extractor using the last layer of ${D_{(\\cdot)}}$, and $\\gamma \\in[0,1]$ indicates the ratio between discriminator CNN feature level and the pixel level loss.\n\nIn the paper, the authors also weight the cycle consistency loss by the quality of the generated images obtained from the discriminator outputs. See the extended formula from the paper (for simplicity in the notebook I use the 1st formula without multiplier $D_X(x)$ ):\n\n\\begin{aligned} & \\tilde{\\mathcal{L}}_{\\mathrm{cyc}}\\left(G, F, D_X, X, \\gamma\\right)= \\mathbb{E}_{x \\sim p_{\\text {data }}(x)}\\left[D_X(x)\\left(\\gamma\\left\\|f_{D_X}(F(G(x)))-f_{D_X}(x)\\right\\|_1+(1-\\gamma)\\|F(G(x))-x\\|_1\\right)\\right]\\end{aligned}\n\nThe full objective at epoch t:\n\n\\begin{aligned}\n\\mathcal{L}\\left(G, F, D_X, D_Y, t\\right) & =\\mathcal{L}_{\\mathrm{GAN}}\\left(G, D_Y, X, Y\\right)+\\mathcal{L}_{\\mathrm{GAN}}\\left(F, D_X, Y, X\\right) \\\\\n& +\\lambda_t \\tilde{\\mathcal{L}}_{\\mathrm{cyc}}\\left(G, F, D_X, X, \\gamma_t\\right)+\\lambda_t \\tilde{\\mathcal{L}}_{\\mathrm{cyc}}\\left(F, G, D_Y, Y, \\gamma_t\\right),\n\\end{aligned}\n\nwhere the authors suggest $\\lambda_t$ to linearly decrease to a small value and $\\gamma_t$ to linearly increase to a value close to 1.\n\nThe authors also train the model with a constant learning rate of 0:0002 for 100 iterations and a linearly decaying learning rate to 0 for another 100 iterations. I use this approach as well.\n\nAnother optimization I use in my notebook is to add noise layers to the discriminator network. This makes the discriminator weaker and prevents it from quickly overfitting (similar to the two-objective discriminator in [UnfrendlyAI's notebook](https://www.kaggle.com/code/unfriendlyai/two-objective-discriminator). DiffAugment is also used.\n\nWhile experimenting the authors of the paper [CycleGAN with Better Cycles](https://www.tongzhouwang.info/better_cycles/report.pdf) found out that the quality of the result is very sensitive to different parameters. I experimented with a lot of combinations, but I can't achive the same results as authors (maybe also because on competition limitation of 300 min). So the main purpose of this task for me is educational.","metadata":{"id":"gBuaqfO9MRfL"}},{"cell_type":"markdown","source":"## Dependencies","metadata":{"id":"XgSPG-vBMRfM"}},{"cell_type":"code","source":"%matplotlib inline\nimport warnings\nwarnings.simplefilter(action='ignore', category=(RuntimeWarning, FutureWarning, UserWarning))\n\nimport gc\nimport time\nimport os\nimport urllib.request\nimport shutil\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom PIL import Image\nimport zipfile\nfrom io import BytesIO\nfrom matplotlib.pyplot import subplots\nimport tensorflow as tf\nimport keras\nfrom keras import layers, Model\nfrom keras.models import Sequential\nfrom keras.layers import Layer, Dense, Activation, Dropout, Input, concatenate, Average, Concatenate, Conv2D, Conv2DTranspose, BatchNormalization, ReLU, LeakyReLU, ZeroPadding2D, GaussianNoise\nfrom keras.optimizers import Adam\nfrom keras.utils import array_to_img\nfrom keras.callbacks import Callback\nfrom keras.initializers import RandomNormal\nfrom keras.losses import BinaryCrossentropy\n\nnp.random.seed(0)\ntf.random.set_seed(0)\n\nIN_COLAB = os.getenv(\"COLAB_RELEASE_TAG\") is not None\nON_KAGGLE = os.getenv(\"KAGGLE_KERNEL_RUN_TYPE\") is not None\n\nAUTOTUNE = tf.data.AUTOTUNE\n\nprint(\"tf version: \", tf.__version__)","metadata":{"id":"KmQY5gdYMRfN","outputId":"345fec3d-894d-4e94-b4ea-16f88d1316cd","_kg_hide-output":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2023-08-06T05:12:20.505252Z","iopub.execute_input":"2023-08-06T05:12:20.506469Z","iopub.status.idle":"2023-08-06T05:12:20.523079Z","shell.execute_reply.started":"2023-08-06T05:12:20.506413Z","shell.execute_reply":"2023-08-06T05:12:20.522033Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Download the <a href=\"https://github.com/mit-han-lab/data-efficient-gans/blob/master/DiffAugment-stylegan2/DiffAugment_tf.py\">DiffAugment_tf.py</a> file from the data-efficient-gans project and import all the functions from it. Alternatively we can copy the code of all these functions to notebook.","metadata":{"id":"d_H-y_lrMRfQ"}},{"cell_type":"code","source":"url = 'https://raw.githubusercontent.com/mit-han-lab/data-efficient-gans/master/DiffAugment-stylegan2/DiffAugment_tf.py'\nfilename = 'DiffAugment_tf.py'\nurllib.request.urlretrieve(url, filename)\nimport DiffAugment_tf\n\nfunction_names = [name for name in dir(DiffAugment_tf) if callable(getattr(DiffAugment_tf, name))]\nprint(\"DiffAugment functions:\\n\", function_names)","metadata":{"id":"ihjI1HZrMRfQ","outputId":"0a939164-3db1-4511-809a-6bc82f3e6686","execution":{"iopub.status.busy":"2023-08-06T05:12:24.059014Z","iopub.execute_input":"2023-08-06T05:12:24.059429Z","iopub.status.idle":"2023-08-06T05:12:24.233940Z","shell.execute_reply.started":"2023-08-06T05:12:24.059393Z","shell.execute_reply":"2023-08-06T05:12:24.233064Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## TPU configuration","metadata":{"id":"HeyboZcqMRfR"}},{"cell_type":"code","source":"if ON_KAGGLE:\n    try:\n        tpu = tf.distribute.cluster_resolver.TPUClusterResolver.connect(tpu=\"local\")\n        strategy = tf.distribute.TPUStrategy(tpu)\n        print(\"on TPU\")\n    except tf.errors.NotFoundError:\n        print(\"not on TPU\")\n        strategy = tf.distribute.MirroredStrategy()\n\n    print(\"REPLICAS: \", strategy.num_replicas_in_sync)","metadata":{"id":"6mXZSWJEMRfR","_kg_hide-output":true,"execution":{"iopub.status.busy":"2023-08-06T05:12:28.856944Z","iopub.execute_input":"2023-08-06T05:12:28.857312Z","iopub.status.idle":"2023-08-06T05:12:36.824992Z","shell.execute_reply.started":"2023-08-06T05:12:28.857281Z","shell.execute_reply":"2023-08-06T05:12:36.823805Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if IN_COLAB:\n    try:\n        tpu = tf.distribute.cluster_resolver.TPUClusterResolver()  # TPU detection\n        print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])\n    except ValueError:\n        raise BaseException('ERROR: Not connected to a TPU runtime; please see the previous cell in this notebook for instructions!')\n\n    tf.config.experimental_connect_to_cluster(tpu)\n    tf.tpu.experimental.initialize_tpu_system(tpu)\n    strategy = tf.distribute.TPUStrategy(tpu)\n    print(\"REPLICAS: \", strategy.num_replicas_in_sync)","metadata":{"execution":{"iopub.status.busy":"2023-08-05T22:44:51.346842Z","iopub.execute_input":"2023-08-05T22:44:51.347153Z","iopub.status.idle":"2023-08-05T22:44:51.353793Z","shell.execute_reply.started":"2023-08-05T22:44:51.347125Z","shell.execute_reply":"2023-08-05T22:44:51.352981Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Data Collection <a class=\"anchor\" id=\"datacollection\"></a>","metadata":{"id":"Mkc0z38zMRfS"}},{"cell_type":"markdown","source":"I sort files they to have same order as from KaggleDatasets().get_gcs_path","metadata":{}},{"cell_type":"code","source":"if ON_KAGGLE:\n    competition = 'gan-getting-started'\n    photo_path_jpg = f\"/kaggle/input/{competition}/photo_jpg\"\n    monet_path_jpg = f\"/kaggle/input/{competition}/monet_jpg\"\n    photo_path_tfrec = f\"/kaggle/input/{competition}/photo_tfrec\"\n    monet_path_tfrec = f\"/kaggle/input/{competition}/monet_tfrec\"\n\n    photo_files_jpg = tf.io.gfile.glob(f\"{photo_path_jpg}/*.jpg\")\n    monet_files_jpg = tf.io.gfile.glob(f\"{monet_path_jpg}/*.jpg\")\n    photo_files_tfrec = sorted(tf.io.gfile.glob(f\"{photo_path_tfrec}/*.tfrec\"))\n    monet_files_tfrec = sorted(tf.io.gfile.glob(f\"{monet_path_tfrec}/*.tfrec\"))","metadata":{"id":"V56246fdMRfS","execution":{"iopub.status.busy":"2023-08-06T05:20:04.681626Z","iopub.execute_input":"2023-08-06T05:20:04.682037Z","iopub.status.idle":"2023-08-06T05:20:05.693217Z","shell.execute_reply.started":"2023-08-06T05:20:04.682002Z","shell.execute_reply":"2023-08-06T05:20:05.691997Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"In Google Colab, datasets can be retrieved from Google Cloud Storage (and it's mandatory for TPU configuration).","metadata":{"id":"_42ruW3yMRfS"}},{"cell_type":"code","source":"# Run in the Kaggle environment (without TPU):\n# from kaggle_datasets import KaggleDatasets\n# gcs_path = KaggleDatasets().get_gcs_path('gan-getting-started')\n# !gsutil ls $gcs_path\n# Replace pathes values below:\nif IN_COLAB:\n    photo_files_tfrec = !gsutil ls 'gs://kds-6352fc4af0e31cf896c4f891fb70b8cfe4c4ac7a83a05f5b4d2bdd71/photo_tfrec/'\n    monet_files_tfrec = !gsutil ls 'gs://kds-6352fc4af0e31cf896c4f891fb70b8cfe4c4ac7a83a05f5b4d2bdd71/monet_tfrec/'","metadata":{"execution":{"iopub.status.busy":"2023-08-05T22:44:52.118366Z","iopub.execute_input":"2023-08-05T22:44:52.118738Z","iopub.status.idle":"2023-08-05T22:44:52.125365Z","shell.execute_reply.started":"2023-08-05T22:44:52.118708Z","shell.execute_reply":"2023-08-05T22:44:52.124415Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"photo_files_tfrds = tf.data.TFRecordDataset(photo_files_tfrec)\nmonet_files_tfrds = tf.data.TFRecordDataset(monet_files_tfrec)","metadata":{"id":"mTQ68FDRMRfU","execution":{"iopub.status.busy":"2023-08-05T22:44:52.126509Z","iopub.execute_input":"2023-08-05T22:44:52.126826Z","iopub.status.idle":"2023-08-05T22:44:52.167306Z","shell.execute_reply.started":"2023-08-05T22:44:52.126799Z","shell.execute_reply":"2023-08-05T22:44:52.166476Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Data Analysis <a class=\"anchor\" id=\"dataanalysis\"></a>","metadata":{"id":"s2z_fI7lMRfU"}},{"cell_type":"code","source":"def count_tfrecord_examples(tfrecords_dir: str,) -> int:\n    \"\"\"\n    Counts the total number of examples in a collection of TFRecord files.\n\n    :param tfrecords_dir: directory that is assumed to contain only TFRecord files\n    :return: the total number of examples in the collection of TFRecord files\n        found in the specified directory\n    \"\"\"\n    count = 0\n    for file_name in os.listdir(tfrecords_dir):\n        tfrecord_path = os.path.join(tfrecords_dir, file_name)\n        count += tf.data.TFRecordDataset(tfrecord_path).reduce(np.int64(0), lambda x, _: x + 1)\n    return count.numpy()","metadata":{"id":"hQUMq2WTMRfU","execution":{"iopub.status.busy":"2023-08-05T22:44:52.168394Z","iopub.execute_input":"2023-08-05T22:44:52.168676Z","iopub.status.idle":"2023-08-05T22:44:52.175306Z","shell.execute_reply.started":"2023-08-05T22:44:52.168652Z","shell.execute_reply":"2023-08-05T22:44:52.174431Z"},"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if ON_KAGGLE:\n    print('Photo .jpg files count:', len(photo_files_jpg))\n    print('Monet .jpg files count:', len(monet_files_jpg))\n\nprint('Photo .tfrec files count:', len(photo_files_tfrec))","metadata":{"id":"x4ivk4gvMRfV","execution":{"iopub.status.busy":"2023-08-05T22:44:52.176525Z","iopub.execute_input":"2023-08-05T22:44:52.176826Z","iopub.status.idle":"2023-08-05T22:44:52.191343Z","shell.execute_reply.started":"2023-08-05T22:44:52.176799Z","shell.execute_reply":"2023-08-05T22:44:52.190418Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Data visualization <a class=\"anchor\" id=\"datavisualization\"></a>","metadata":{"id":"N9hgSJURMRfV"}},{"cell_type":"markdown","source":"### Visualize .jpg images","metadata":{"id":"-JOZ8scMMRfV"}},{"cell_type":"code","source":"def plot_image(image):\n    plt.imshow(image)\n    plt.axis('off')\n    plt.show()\n    \ndef open_and_plot_image(image_path):\n    image = Image.open(image_path)\n    plot_image(image)\n    \ndef plot_tfrec_image(record, image_feature, extra_features=[]):\n    example = tf.train.Example()\n    example.ParseFromString(record.numpy())\n    image = tf.image.decode_jpeg(\n        example.features.feature[image_feature].bytes_list.value[0],\n        channels=3\n    )\n    for feature in extra_features:\n        print('Image type:\\t', example.features.feature[feature].bytes_list.value[0].decode('utf-8'))\n    plot_image(image)","metadata":{"id":"I8QgGmGwMRfW","execution":{"iopub.status.busy":"2023-08-06T05:20:52.446918Z","iopub.execute_input":"2023-08-06T05:20:52.447292Z","iopub.status.idle":"2023-08-06T05:20:52.455819Z","shell.execute_reply.started":"2023-08-06T05:20:52.447261Z","shell.execute_reply":"2023-08-06T05:20:52.454701Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if ON_KAGGLE:\n    print('Photo:')\n    open_and_plot_image(os.path.join(photo_path_jpg, photo_files_jpg[0]))","metadata":{"id":"NgwbwVwZMRfW","execution":{"iopub.status.busy":"2023-08-05T22:44:52.209350Z","iopub.execute_input":"2023-08-05T22:44:52.209854Z","iopub.status.idle":"2023-08-05T22:44:52.463398Z","shell.execute_reply.started":"2023-08-05T22:44:52.209825Z","shell.execute_reply":"2023-08-05T22:44:52.462475Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if ON_KAGGLE:\n    print('Monet:')\n    open_and_plot_image(os.path.join(monet_path_jpg, monet_files_jpg[0]))","metadata":{"id":"Re9Cu4sKMRfX","execution":{"iopub.status.busy":"2023-08-05T22:44:52.464554Z","iopub.execute_input":"2023-08-05T22:44:52.464859Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Visualize .tfrec images","metadata":{"id":"mgz8v2YpMRfX"}},{"cell_type":"markdown","source":"Get the structure of the TFRecord sample protobuffer.","metadata":{"id":"V1XstSRqMRfY"}},{"cell_type":"code","source":"def get_tfrecord_structure(ds):\n    record = next(iter(ds.take(1)))\n    example = tf.train.Example()\n    example.ParseFromString(record.numpy())\n    print('TFRecord structure:')\n    for feature_key in example.features.feature.keys():\n        print(feature_key)","metadata":{"id":"MUYEDptgMRfY","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"get_tfrecord_structure(monet_files_tfrds)","metadata":{"id":"fQuW1SpPMRfY","outputId":"9ee24752-fcbd-4c4e-d686-3726a9ba5515","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"get_tfrecord_structure(photo_files_tfrds)","metadata":{"id":"YzZKcF8gMRfY","outputId":"704ef3c5-d780-4664-ecb0-47110d0ddfda","execution":{"iopub.execute_input":"2023-08-05T22:44:52.672385Z","iopub.status.idle":"2023-08-05T22:44:52.702932Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_tfrec_image(\n    next(iter(photo_files_tfrds.skip(np.random.randint(0,7038)))),\n    'image', ['target', 'image_name']\n)","metadata":{"id":"MkT1l1RYMRfZ","outputId":"d331c4e2-21c1-4830-f93d-931163c1982e","execution":{"iopub.status.busy":"2023-08-05T22:44:52.704130Z","iopub.execute_input":"2023-08-05T22:44:52.704558Z","iopub.status.idle":"2023-08-05T22:44:56.442292Z","shell.execute_reply.started":"2023-08-05T22:44:52.704527Z","shell.execute_reply":"2023-08-05T22:44:56.441250Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_tfrec_image(\n    next(iter(monet_files_tfrds.skip(np.random.randint(0,300)))),\n    'image', ['target', 'image_name']\n)","metadata":{"id":"7GKS8VclMRfZ","outputId":"659985db-28a4-41be-f3ff-c8a658eacef9","execution":{"iopub.status.busy":"2023-08-05T22:44:56.443426Z","iopub.execute_input":"2023-08-05T22:44:56.443800Z","iopub.status.idle":"2023-08-05T22:44:56.730018Z","shell.execute_reply.started":"2023-08-05T22:44:56.443774Z","shell.execute_reply":"2023-08-05T22:44:56.729111Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Image names corresponds to the same images in *_jpg folders.","metadata":{"id":"1n3AwyJXMRfa"}},{"cell_type":"markdown","source":"## EDA Summary <a class=\"anchor\" id=\"edasummary\"></a>","metadata":{"id":"q0DyTSbKMRfa"}},{"cell_type":"markdown","source":"1. Photos dataset has 7038 items\n1. Monet dataset has 300 items\n1. Image size: 256x256, channels: 3","metadata":{"id":"lbC2D_N7MRfa"}},{"cell_type":"markdown","source":"# Data Preprocessing <a class=\"anchor\" id=\"datapreprocessing\"></a>","metadata":{"id":"EFAfssE5MRfa"}},{"cell_type":"code","source":"if strategy.num_replicas_in_sync == 1:\n    BATCH_SIZE = 2\n    BUFFER_SIZE = 256\nelse:\n    BATCH_SIZE_PER_REPLICA = 1\n    BATCH_SIZE = BATCH_SIZE_PER_REPLICA * strategy.num_replicas_in_sync\n    TEST_BATCH_SIZE_PER_REPLICA = 32\n    TEST_BATCH_SIZE = TEST_BATCH_SIZE_PER_REPLICA * strategy.num_replicas_in_sync","metadata":{"id":"a3x55SuwMRfb","execution":{"iopub.status.busy":"2023-08-06T05:17:35.839935Z","iopub.execute_input":"2023-08-06T05:17:35.840231Z","iopub.status.idle":"2023-08-06T05:17:35.857016Z","shell.execute_reply.started":"2023-08-06T05:17:35.840202Z","shell.execute_reply":"2023-08-06T05:17:35.855969Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Create functions for jpeg decoding and data augmentation.","metadata":{"id":"WkPz3NxWMRfb"}},{"cell_type":"code","source":"def decode_image(image):\n    shape = [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, shape)\n    return image\n\ndef read_tfrecord(example):\n    tfrecord_format = {\n        \"image\": 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):\n    dataset = tf.data.TFRecordDataset(filenames)\n    dataset = dataset.map(read_tfrecord, num_parallel_calls=AUTOTUNE)\n    return dataset\n\ndef data_augment(image):\n    image = tf.image.resize(image, [286, 286])\n    image = tf.image.random_crop(image, size=[BATCH_SIZE, 256, 256, 3])\n    image = tf.image.random_flip_left_right(image)\n    return image","metadata":{"id":"twKzxNMPMRfb","execution":{"iopub.status.busy":"2023-08-06T05:17:37.618682Z","iopub.execute_input":"2023-08-06T05:17:37.619744Z","iopub.status.idle":"2023-08-06T05:17:37.629920Z","shell.execute_reply.started":"2023-08-06T05:17:37.619705Z","shell.execute_reply":"2023-08-06T05:17:37.628673Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def prepare_gan_ds(photo_files, monet_files, augment_fn=None, repeat=True, shuffle=True, cache=True, batch_size=1):\n    monet_ds = load_dataset(monet_files)\n    photo_ds = load_dataset(photo_files)\n    if cache:\n        monet_ds = monet_ds.cache()\n        photo_ds = photo_ds.cache()\n    if repeat:\n        monet_ds = monet_ds.repeat()\n        photo_ds = photo_ds.repeat()\n    if shuffle:\n        monet_ds_entries = monet_ds.reduce(0, lambda x,_: x+1).numpy()\n        photo_ds_entries = photo_ds.reduce(0, lambda x,_: x+1).numpy()\n        monet_ds = monet_ds.shuffle(monet_ds_entries)\n        photo_ds = photo_ds.shuffle(photo_ds_entries)\n    monet_ds = monet_ds.batch(batch_size)\n    photo_ds = photo_ds.batch(batch_size)\n    if augment_fn:\n        monet_ds = monet_ds.map(augment_fn, num_parallel_calls=AUTOTUNE)\n        photo_ds = photo_ds.map(augment_fn, num_parallel_calls=AUTOTUNE)\n    monet_ds = monet_ds.prefetch(AUTOTUNE)\n    photo_ds = photo_ds.prefetch(AUTOTUNE)\n    gan_ds = tf.data.Dataset.zip((photo_ds, monet_ds))\n    return gan_ds","metadata":{"id":"2pYptugtMRf1","execution":{"iopub.status.busy":"2023-08-06T05:17:39.363651Z","iopub.execute_input":"2023-08-06T05:17:39.364680Z","iopub.status.idle":"2023-08-06T05:17:39.375896Z","shell.execute_reply.started":"2023-08-06T05:17:39.364638Z","shell.execute_reply":"2023-08-06T05:17:39.374582Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_ds = prepare_gan_ds(\n    photo_files_tfrec,\n    monet_files_tfrec,\n    augment_fn=data_augment,\n    repeat=True,\n    shuffle=False,\n    cache=False,\n    batch_size=BATCH_SIZE\n)\ntest_photo_ds = load_dataset(photo_files_tfrec).batch(TEST_BATCH_SIZE).prefetch(TEST_BATCH_SIZE)\ntrain_ds.element_spec, test_photo_ds.element_spec","metadata":{"id":"k-D7sA5WMRf2","outputId":"116287e5-5966-42ac-fe66-9531656a1507","execution":{"iopub.status.busy":"2023-08-06T05:20:40.768791Z","iopub.execute_input":"2023-08-06T05:20:40.769602Z","iopub.status.idle":"2023-08-06T05:20:40.924209Z","shell.execute_reply.started":"2023-08-06T05:20:40.769563Z","shell.execute_reply":"2023-08-06T05:20:40.922920Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Check datasets.","metadata":{}},{"cell_type":"code","source":"random_photo_image, random_monet_image = next(iter(train_ds.take(1)))\nrandom_photo_image_test = next(iter(test_photo_ds.take(1)))\nprint('Photo image pixel range: min: {0:.2f}, max: {1:.2f}'.format(random_photo_image.numpy().min(), random_photo_image.numpy().max()))\nprint('Monet image pixel range: min: {0:.2f}, max: {1:.2f}'.format(random_monet_image.numpy().min(), random_monet_image.numpy().max()))\nprint('Photo (test) image pixel range: min: {0:.2f}, max: {1:.2f}'.format(random_photo_image_test.numpy().min(), random_photo_image_test.numpy().max()))","metadata":{"id":"lNOlcNS4MRf2","outputId":"91549ef8-bb3a-4874-d008-00713839d9b4","execution":{"iopub.status.busy":"2023-08-06T05:20:42.711192Z","iopub.execute_input":"2023-08-06T05:20:42.712285Z","iopub.status.idle":"2023-08-06T05:20:43.295074Z","shell.execute_reply.started":"2023-08-06T05:20:42.712243Z","shell.execute_reply":"2023-08-06T05:20:43.293884Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_image((random_photo_image[0]+1)/2)","metadata":{"id":"fAWNZZsyMRf3","outputId":"322936f7-2f14-425e-ec50-ccac4fd2b00d","execution":{"iopub.status.busy":"2023-08-06T05:20:58.096498Z","iopub.execute_input":"2023-08-06T05:20:58.097589Z","iopub.status.idle":"2023-08-06T05:20:58.351489Z","shell.execute_reply.started":"2023-08-06T05:20:58.097548Z","shell.execute_reply":"2023-08-06T05:20:58.350255Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_image((random_monet_image[0]+1)/2)","metadata":{"id":"DvCEb2LQMRf3","outputId":"60acb728-de70-4964-8b33-b82edb5a759e","execution":{"iopub.status.busy":"2023-08-05T22:44:57.797489Z","iopub.execute_input":"2023-08-05T22:44:57.797777Z","iopub.status.idle":"2023-08-05T22:44:57.945693Z","shell.execute_reply.started":"2023-08-05T22:44:57.797752Z","shell.execute_reply":"2023-08-05T22:44:57.944728Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_image((random_photo_image_test[0]+1)/2)","metadata":{"id":"atFPO6HFMRf4","outputId":"50bfa33b-50f6-40d5-953a-365b02f3740e","execution":{"iopub.status.busy":"2023-08-05T22:44:57.946844Z","iopub.execute_input":"2023-08-05T22:44:57.947143Z","iopub.status.idle":"2023-08-05T22:44:58.090516Z","shell.execute_reply.started":"2023-08-05T22:44:57.947117Z","shell.execute_reply":"2023-08-05T22:44:58.089293Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"del random_photo_image, random_monet_image, random_photo_image_test\ngc.collect()","metadata":{"id":"neWcKLcZMRf4","outputId":"48a5890f-64d4-4a4a-9d83-453aaed54fcf","execution":{"iopub.status.busy":"2023-08-05T22:44:58.091882Z","iopub.execute_input":"2023-08-05T22:44:58.092313Z","iopub.status.idle":"2023-08-05T22:44:58.306158Z","shell.execute_reply.started":"2023-08-05T22:44:58.092283Z","shell.execute_reply":"2023-08-05T22:44:58.305271Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model Training and Evaluation <a class=\"anchor\" id=\"trainandevaluate\"></a>","metadata":{"id":"XwFmmRd5MRf5"}},{"cell_type":"markdown","source":"First, create the necessary classes, functions, and objects.","metadata":{"id":"yAu_w8KUMRf5"}},{"cell_type":"code","source":"def diffaug_fn(image):\n    return DiffAugment_tf.DiffAugment(image, \"color,translation,cutout\")","metadata":{"id":"txfVfx1wMRf5","execution":{"iopub.status.busy":"2023-08-05T22:44:58.307263Z","iopub.execute_input":"2023-08-05T22:44:58.307568Z","iopub.status.idle":"2023-08-05T22:44:58.311849Z","shell.execute_reply.started":"2023-08-05T22:44:58.307541Z","shell.execute_reply":"2023-08-05T22:44:58.310920Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Generator, Discriminator <a class=\"anchor\" id=\"gendisc\"></a>","metadata":{"id":"pNMhkCUoMRf6"}},{"cell_type":"code","source":"class InstanceNormalization(tf.keras.layers.Layer):\n    def __init__(self, epsilon=1e-5):\n        super(InstanceNormalization, self).__init__()\n        self.epsilon = epsilon\n\n    def build(self, input_shape):\n        self.scale = self.add_weight(\n            name='scale',\n            shape=input_shape[-1:],\n            initializer=tf.random_normal_initializer(1., 0.02),\n            trainable=True)\n        self.offset = self.add_weight(\n            name='offset',\n            shape=input_shape[-1:],\n            initializer='zeros',\n            trainable=True)\n\n    def call(self, x):\n        mean, variance = tf.nn.moments(x, axes=[1, 2], keepdims=True)\n        inv = tf.math.rsqrt(variance + self.epsilon)\n        normalized = (x - mean) * inv\n        return self.scale * normalized + self.offset","metadata":{"id":"pFdFcLtmMRf6","execution":{"iopub.status.busy":"2023-08-05T22:44:58.318216Z","iopub.execute_input":"2023-08-05T22:44:58.318582Z","iopub.status.idle":"2023-08-05T22:44:58.328593Z","shell.execute_reply.started":"2023-08-05T22:44:58.318555Z","shell.execute_reply":"2023-08-05T22:44:58.327569Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Create functions to construct the generator and discriminator.","metadata":{"id":"T9YLevs0MRf6"}},{"cell_type":"code","source":"def downsample(filters, size, apply_instancenorm=True, add_noise=False):\n    initializer = tf.random_normal_initializer(0., 0.02)\n    result = keras.Sequential()\n    result.add(Conv2D(filters, size, strides=2, padding='same',\n                      kernel_initializer=initializer, use_bias=False))\n    if add_noise:\n        result.add(GaussianNoise(0.2))\n    if apply_instancenorm:\n        result.add(InstanceNormalization())\n    result.add(LeakyReLU())\n    return result\n\ndef upsample(filters, size, apply_dropout=False):\n    initializer = tf.random_normal_initializer(0., 0.02)\n    result = Sequential()\n    result.add(Conv2DTranspose(filters, size, strides=2,\n                               padding='same',\n                               kernel_initializer=initializer,\n                               use_bias=False))\n    result.add(InstanceNormalization())\n    if apply_dropout:\n        result.add(Dropout(0.5))\n    result.add(ReLU())\n    return result","metadata":{"execution":{"iopub.status.busy":"2023-08-05T22:44:58.329619Z","iopub.execute_input":"2023-08-05T22:44:58.329885Z","iopub.status.idle":"2023-08-05T22:44:58.343120Z","shell.execute_reply.started":"2023-08-05T22:44:58.329861Z","shell.execute_reply":"2023-08-05T22:44:58.342261Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def create_generator():\n    inputs = Input(shape=[256,256,3])\n    output_channels = 3\n\n    down_stack = [\n        downsample(64, 4, apply_instancenorm=False), # (bs, 128, 128, 64)\n        downsample(128, 4), # (bs, 64, 64, 128)\n        downsample(256, 4), # (bs, 32, 32, 256)\n        downsample(512, 4), # (bs, 16, 16, 512)\n        downsample(512, 4), # (bs, 8, 8, 512)\n        downsample(512, 4), # (bs, 4, 4, 512)\n        downsample(512, 4), # (bs, 2, 2, 512)\n        downsample(512, 4), # (bs, 1, 1, 512)\n    ]\n\n    up_stack = [\n        upsample(512, 4, apply_dropout=True), # (bs, 2, 2, 1024)\n        upsample(512, 4, apply_dropout=True), # (bs, 4, 4, 1024)\n        upsample(512, 4, apply_dropout=True), # (bs, 8, 8, 1024)\n        upsample(512, 4), # (bs, 16, 16, 1024)\n        upsample(256, 4), # (bs, 32, 32, 512)\n        upsample(128, 4), # (bs, 64, 64, 256)\n        upsample(64, 4), # (bs, 128, 128, 128)\n    ]\n\n    initializer = tf.random_normal_initializer(0., 0.02)\n    last = Conv2DTranspose(output_channels, 4,\n                           strides=2,\n                           padding='same',\n                           kernel_initializer=initializer,\n                           activation='tanh') # (bs, 256, 256, 3)\n\n    x = inputs\n\n    # Downsampling through the model\n    skips = []\n    for down in down_stack:\n        x = down(x)\n        skips.append(x)\n\n    skips = reversed(skips[:-1])\n\n    # Upsampling and establishing the skip connections\n    for up, skip in zip(up_stack, skips):\n        x = up(x)\n        x = Concatenate()([x, skip])\n    outputs = last(x)\n    return keras.Model(inputs=inputs, outputs=outputs)","metadata":{"execution":{"iopub.status.busy":"2023-08-05T22:44:58.344090Z","iopub.execute_input":"2023-08-05T22:44:58.344374Z","iopub.status.idle":"2023-08-05T22:44:58.356987Z","shell.execute_reply.started":"2023-08-05T22:44:58.344349Z","shell.execute_reply":"2023-08-05T22:44:58.356200Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def create_discriminator(add_noise=True):\n    initializer = tf.random_normal_initializer(0., 0.02)\n    inputs = Input(shape=[256, 256, 3], name='input_image')\n    x = inputs\n    x = downsample(64, 4, False, add_noise=add_noise)(x) # (bs, 128, 128, 64)\n    x = downsample(128, 4, add_noise=add_noise)(x) # (bs, 64, 64, 128)\n    x = downsample(256, 4, add_noise=add_noise)(x) # (bs, 32, 32, 256)\n    x = ZeroPadding2D()(x) # (bs, 34, 34, 256)\n    x = Conv2D(512, 4, strides=1,\n               kernel_initializer=initializer,\n               use_bias=False)(x) # (bs, 31, 31, 512)\n    if add_noise:\n        x = GaussianNoise(0.2)(x)\n    x = InstanceNormalization()(x)\n    x = LeakyReLU()(x)\n    x = ZeroPadding2D()(x) # (bs, 33, 33, 512)\n    outputs = Conv2D(1, 4, strides=1,\n                     kernel_initializer=initializer)(x) # (bs, 30, 30, 1)\n    return tf.keras.Model(inputs=inputs, outputs=outputs)","metadata":{"execution":{"iopub.status.busy":"2023-08-05T22:44:58.358026Z","iopub.execute_input":"2023-08-05T22:44:58.358327Z","iopub.status.idle":"2023-08-05T22:44:58.373209Z","shell.execute_reply.started":"2023-08-05T22:44:58.358300Z","shell.execute_reply":"2023-08-05T22:44:58.372298Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model <a class=\"anchor\" id=\"model\"></a>","metadata":{}},{"cell_type":"markdown","source":"Implement model training class with losses methods and redefined train_step and compile methods. The cycle_loss_x and cycle_loss_y variables in the train_step method contain the L1 feature matching loss in addition to the standard CycleGAN cycle loss.","metadata":{"id":"-dLO6mmfMRgA"}},{"cell_type":"code","source":"class TrainingGAN(Model):\n    def __init__(self, generator_g, discriminator_x, generator_f, discriminator_y,\n                 lambda_loss=10, gamma_loss=1e-4, lambda_id_loss=1e-5, diffaug_fn=None, batch_size=32, **kwargs):\n        super().__init__(**kwargs)\n        self.generator_g = generator_g\n        self.generator_f = generator_f\n        self.discriminator_x = discriminator_x\n        self.discriminator_y = discriminator_y\n        self.lambda_loss = lambda_loss\n        self.gamma_loss = gamma_loss\n        self.lambda_id_loss = lambda_id_loss\n        self.diffaug_fn = diffaug_fn\n        disc_inputs = Input(shape=[None, None, 3], name='input_image')\n        self.discriminator_features_x = Model(inputs=self.discriminator_x.input,\n                                              outputs=self.discriminator_x.layers[-2].output)\n        self.discriminator_features_y = Model(inputs=self.discriminator_y.input,\n                                              outputs=self.discriminator_y.layers[-2].output)\n\n    def compile(self, generator_g_optimizer, discriminator_x_optimizer,\n                generator_f_optimizer, discriminator_y_optimizer):\n        super().compile()\n        self.generator_g_optimizer = generator_g_optimizer\n        self.discriminator_x_optimizer = discriminator_x_optimizer\n        self.generator_f_optimizer = generator_f_optimizer\n        self.discriminator_y_optimizer = discriminator_y_optimizer\n\n    def _discriminator_bce_loss(self, real, generated):\n        real_loss = BinaryCrossentropy(\n            from_logits=True,\n            reduction=tf.keras.losses.Reduction.NONE)(tf.ones_like(real), real)\n        real_loss = tf.reduce_mean(real_loss)\n        generated_loss = BinaryCrossentropy(\n            from_logits=True,\n            reduction=tf.keras.losses.Reduction.NONE)(tf.zeros_like(generated), generated)\n        generated_loss = tf.reduce_mean(generated_loss)\n        total_disc_loss = real_loss + generated_loss\n        return total_disc_loss * 0.5\n\n    def _generator_bce_loss(self, generated):\n        loss = BinaryCrossentropy(\n            from_logits=True,\n            reduction=tf.keras.losses.Reduction.NONE)(tf.ones_like(generated), generated)\n        loss = tf.reduce_mean(loss)\n        return loss\n\n    def _cycle_loss(self, real_image, cycled_image):\n        loss = tf.reduce_mean(tf.abs(real_image - cycled_image))\n        return loss\n\n    def _identity_loss(self, real_image, same_image):\n        loss = tf.reduce_mean(tf.abs(real_image - same_image))\n        return 0.5 * loss\n\n    @tf.function\n    def train_step(self, batch):\n        (real_x, real_y) = batch\n        batch_size = tf.shape(real_y)[0]\n        with tf.GradientTape(persistent=True) as tape:\n            fake_y = self.generator_g(real_x, training=True)\n            cycled_x = self.generator_f(fake_y, training=True)\n            fake_x = self.generator_f(real_y, training=True)\n            cycled_y = self.generator_g(fake_x, training=True)\n            same_x = self.generator_f(real_x, training=True)\n            same_y = self.generator_g(real_y, training=True)\n\n            if self.diffaug_fn:\n                both_y = tf.concat([real_y, fake_y], axis=0)\n                aug_y = self.diffaug_fn(both_y)\n                aug_real_y = aug_y[:batch_size]\n                aug_fake_y = aug_y[batch_size:]\n                disc_real_y = self.discriminator_y(aug_real_y, training=True)\n                disc_fake_y = self.discriminator_y(aug_fake_y, training=True)\n            else:\n                disc_real_y = self.discriminator_y(real_y, training=True)\n                disc_fake_y = self.discriminator_y(fake_y, training=True)\n\n            disc_real_x = self.discriminator_x(real_x, training=True)\n            disc_fake_x = self.discriminator_x(fake_x, training=True)\n\n            disc_feat_x = self.discriminator_features_x(real_x, training=True)\n            disc_feat_cycled_x = self.discriminator_features_x(cycled_x, training=True)\n            disc_feat_y = self.discriminator_features_x(real_y, training=True)\n            disc_feat_cycled_y = self.discriminator_features_x(cycled_y, training=True)\n\n            gen_g_loss = self._generator_bce_loss(disc_fake_y)\n            gen_f_loss = self._generator_bce_loss(disc_fake_x)\n            cycle_loss_x = ((1 - self.gamma_loss) * self._cycle_loss(real_x, cycled_x) +\n                            self.gamma_loss * self._cycle_loss(disc_feat_x, disc_feat_cycled_x))\n            cycle_loss_y = ((1 - self.gamma_loss) *  self._cycle_loss(real_y, cycled_y) +\n                            self.gamma_loss * self._cycle_loss(disc_feat_y, disc_feat_cycled_y))\n            total_cycle_loss = self.lambda_loss * (cycle_loss_x + cycle_loss_y)\n            id_loss_y = self.lambda_id_loss * self._identity_loss(real_y, same_y)\n            id_loss_x = self.lambda_id_loss * self._identity_loss(real_x, same_x)\n            total_gen_g_loss = gen_g_loss + total_cycle_loss + id_loss_y\n            total_gen_f_loss = gen_f_loss + total_cycle_loss + id_loss_x\n\n            disc_x_loss = self._discriminator_bce_loss(disc_real_x, disc_fake_x)\n            disc_y_loss = self._discriminator_bce_loss(disc_real_y, disc_fake_y)\n\n        generator_g_gradients = tape.gradient(total_gen_g_loss, self.generator_g.trainable_variables)\n        generator_f_gradients = tape.gradient(total_gen_f_loss, self.generator_f.trainable_variables)\n        discriminator_x_gradients = tape.gradient(disc_x_loss, self.discriminator_x.trainable_variables)\n        discriminator_y_gradients = tape.gradient(disc_y_loss, self.discriminator_y.trainable_variables)\n\n        self.generator_g_optimizer.apply_gradients(zip(generator_g_gradients, self.generator_g.trainable_variables))\n        self.generator_f_optimizer.apply_gradients(zip(generator_f_gradients, self.generator_f.trainable_variables))\n        self.discriminator_x_optimizer.apply_gradients(zip(discriminator_x_gradients, self.discriminator_x.trainable_variables))\n        self.discriminator_y_optimizer.apply_gradients(zip(discriminator_y_gradients, self.discriminator_y.trainable_variables))\n\n        return {\n            \"disc_x_loss\": disc_x_loss,\n            \"disc_y_loss\": disc_y_loss,\n            \"total_gen_g_loss\": total_gen_g_loss,\n            \"total_gen_f_loss\": total_gen_f_loss,\n            \"gen_g_loss\": cycle_loss_x,\n            \"gen_f_loss\": cycle_loss_y,\n            }","metadata":{"execution":{"iopub.status.busy":"2023-08-05T22:44:58.374268Z","iopub.execute_input":"2023-08-05T22:44:58.374555Z","iopub.status.idle":"2023-08-05T22:44:58.408417Z","shell.execute_reply.started":"2023-08-05T22:44:58.374530Z","shell.execute_reply":"2023-08-05T22:44:58.407562Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Callbacks <a class=\"anchor\" id=\"callbacks\"></a>","metadata":{"id":"E4GBmFKAMRgB"}},{"cell_type":"markdown","source":"Create a training callback that outputs input and predicted images at each training epoch interval.","metadata":{"id":"M01c0aICMRgB"}},{"cell_type":"code","source":"class PlotPredictionsCallback(keras.callbacks.Callback):\n    def __init__(self, input_image, model_generator, epoch_interval=None, nrows=1, figsize=(11, 11)):\n        self.input_image = input_image\n        self.model_generator = model_generator\n        self.epoch_interval = epoch_interval\n        self.nrows = nrows\n        self.figsize=(11, 11)\n\n    def _plot_test_and_pred(self):\n        preds = self.model_generator.predict(self.input_image, verbose=0)\n        (fig, axes) = subplots(nrows=self.nrows, ncols=2, figsize=self.figsize)\n        if self.nrows == 1:\n            axes = [axes]\n        for (ax, inp, pred) in zip(axes, self.input_image, preds):\n            ax[0].imshow(array_to_img(inp))\n            ax[0].set_title(\"Input Image\")\n            ax[0].set_axis_off()\n            ax[1].imshow(array_to_img(pred))\n            ax[1].set_title(\"Prediction\")\n            ax[1].set_axis_off()\n        plt.show()\n\n    def on_epoch_end(self, epoch, logs=None):\n        if self.epoch_interval and epoch % self.epoch_interval == 0:\n            self._plot_test_and_pred()\n\n    def on_train_end(self, logs=None):\n        self._plot_test_and_pred()","metadata":{"id":"82rM_VCnMRgC","execution":{"iopub.status.busy":"2023-08-05T22:44:58.409483Z","iopub.execute_input":"2023-08-05T22:44:58.409760Z","iopub.status.idle":"2023-08-05T22:44:58.426332Z","shell.execute_reply.started":"2023-08-05T22:44:58.409735Z","shell.execute_reply":"2023-08-05T22:44:58.425556Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"FID can be used as a relative training score. Below is a class and a function to compute it. This calculator can also be used in the callback to evaluate the model during training, but it runs relatively slow (about 4 min), so I add this estimation after the model is trained.","metadata":{"id":"ho5-xlsCMRgD"}},{"cell_type":"code","source":"class FIDCalculator(object):\n    def __init__(self, images_x_ds, images_y_ds, model_generator, fid_model_base):\n        self.images_x_ds = images_x_ds\n        self.images_y_ds = images_y_ds\n        self.model_generator = model_generator\n        self.fid_model_base = fid_model_base\n        self.initialized = False\n        self.history = []\n\n    def init_stat_x(self):\n        self.mu_2, self.sigma_2 = self._calculate_activation_statistics_mod(self.images_y_ds, self.fid_model_base)\n        self.initialized = True\n\n    def _calculate_activation_statistics_mod(self, images, fid_model):\n        act = tf.cast(fid_model.predict(images, verbose=0), tf.float32)\n        mu = tf.reduce_mean(act, axis=0)\n        mean_x = tf.reduce_mean(act, axis=0, keepdims=True)\n        mx = tf.matmul(tf.transpose(mean_x), mean_x)\n        vx = tf.matmul(tf.transpose(act), act)/tf.cast(tf.shape(act)[0], tf.float32)\n        sigma = vx - mx\n        return mu, sigma\n\n    def _calculate_frechet_distance(self, mu_1, sigma_1, mu_2, sigma_2):\n        fid_epsilon = 1e-14\n        covmean = tf.linalg.sqrtm(tf.cast(tf.matmul(sigma_1, sigma_2), tf.complex64))\n        covmean = tf.cast(tf.math.real(covmean), tf.float32)\n        tr_covmean = tf.linalg.trace(covmean)\n        fid_value = tf.matmul(\n            tf.expand_dims(mu_1 - mu_2, axis=0),\n            tf.expand_dims(mu_1 - mu_2, axis=1)\n            ) + tf.linalg.trace(sigma_1) + tf.linalg.trace(sigma_2) - 2 * tr_covmean\n        return fid_value\n\n    def _get_gen_plus_fid_model(self):\n        inputs = layers.Input(shape=[256, 256, 3], name='input_image')\n        x = self.model_generator(inputs)\n        outputs = self.fid_model_base(x)\n        fid_model = tf.keras.Model(inputs=inputs, outputs=outputs)\n        return fid_model\n\n    def calc_fid(self):\n        if not self.initialized:\n            self.init_stat_x()\n        fid_model_plus = self._get_gen_plus_fid_model()\n        mu_1, sigma_1 = self._calculate_activation_statistics_mod(self.images_x_ds, fid_model_plus)\n        fid_value = self._calculate_frechet_distance(mu_1, sigma_1, self.mu_2, self.sigma_2)\n        return fid_value","metadata":{"id":"z4ZmwkIPMRgD","execution":{"iopub.status.busy":"2023-08-05T22:44:58.427369Z","iopub.execute_input":"2023-08-05T22:44:58.427659Z","iopub.status.idle":"2023-08-05T22:44:58.444506Z","shell.execute_reply.started":"2023-08-05T22:44:58.427634Z","shell.execute_reply":"2023-08-05T22:44:58.443724Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def create_fid_inception_model():\n    inception_model_base = tf.keras.applications.InceptionV3(\n        input_shape=(256,256,3),\n        pooling=\"avg\",\n        include_top=False)\n    mix3  = inception_model_base.get_layer(\"mixed9\").output\n    f0 = tf.keras.layers.GlobalAveragePooling2D()(mix3)\n    inception_model = tf.keras.Model(inputs=inception_model_base.input, outputs=f0)\n    inception_model.trainable = False\n    return inception_model","metadata":{"id":"kv8Is0BYMRgD","execution":{"iopub.status.busy":"2023-08-05T22:44:58.445519Z","iopub.execute_input":"2023-08-05T22:44:58.445793Z","iopub.status.idle":"2023-08-05T22:44:58.461374Z","shell.execute_reply.started":"2023-08-05T22:44:58.445769Z","shell.execute_reply":"2023-08-05T22:44:58.460475Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class FIDCallback(Callback):\n    def __init__(self, fid_calculator, epoch_interval=None):\n        self.fid_calculator = fid_calculator\n        self.epoch_interval = epoch_interval\n\n    def _get_fid(self):\n        fid = self.fid_calculator.calc_fid()\n        print(\"FID score:\", fid.numpy()[0,0])\n\n    def on_epoch_end(self, epoch, logs=None):\n        if self.epoch_interval and epoch % self.epoch_interval == 0:\n            self._get_fid()\n\n    def on_train_end(self, logs=None):\n        self._get_fid()","metadata":{"id":"WuxICgo4MRgE","execution":{"iopub.status.busy":"2023-08-05T22:44:58.462538Z","iopub.execute_input":"2023-08-05T22:44:58.462850Z","iopub.status.idle":"2023-08-05T22:44:58.475014Z","shell.execute_reply.started":"2023-08-05T22:44:58.462823Z","shell.execute_reply":"2023-08-05T22:44:58.474220Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The next callback updates the learning rates of each model optimizer.","metadata":{}},{"cell_type":"code","source":"class UpdateLearningRateCallback(Callback):\n    def __init__(self, epochs_count, lr_start=2e-4, lr_end=5e-6):\n        super().__init__()\n        self.epoch_min = epochs_count // 2\n        epochs_update_count = epochs_count - self.epoch_min\n        self.lr_values = np.linspace(lr_start, lr_end, epochs_update_count)\n\n    def _scheduler_fn(self, epoch, lr):\n        if epoch < self.epoch_min:\n            return lr\n        else:\n            return self.lr_values[epoch-self.epoch_min]\n\n    def on_epoch_begin(self, epoch, logs=None):\n        self.model.generator_g_optimizer.lr = self._scheduler_fn(epoch, self.model.generator_g_optimizer.lr)\n        self.model.discriminator_x_optimizer.lr = self._scheduler_fn(epoch, self.model.generator_g_optimizer.lr)\n        self.model.generator_f_optimizer.lr = self._scheduler_fn(epoch, self.model.generator_g_optimizer.lr)\n        self.model.discriminator_y_optimizer.lr = self._scheduler_fn(epoch, self.model.generator_g_optimizer.lr)","metadata":{"execution":{"iopub.status.busy":"2023-08-05T22:44:58.476006Z","iopub.execute_input":"2023-08-05T22:44:58.476285Z","iopub.status.idle":"2023-08-05T22:44:58.492375Z","shell.execute_reply.started":"2023-08-05T22:44:58.476261Z","shell.execute_reply":"2023-08-05T22:44:58.491466Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The callback below implements one of the main concepts of the paper [CycleGAN with Better Cycles](https://www.tongzhouwang.info/better_cycles/report.pdf) by updating the lambda and gamma weights during training.","metadata":{}},{"cell_type":"code","source":"class UpdateLossWeightsCallback(Callback):\n    def __init__(self, epochs, lambda_start=10, lambda_end=1e-4, gamma_start=1e-4, gamma_end=0.999):\n        super().__init__()\n        self.epochs = epochs\n        self.lambda_start = lambda_start\n        self.lambda_end = lambda_end\n        self.gamma_start = gamma_start\n        self.gamma_end = gamma_end\n\n    def on_train_begin(self, logs=None):\n        self.lambda_values = np.linspace(self.lambda_start, self.lambda_end, self.epochs)\n        self.gamma_values = np.linspace(self.gamma_start, self.gamma_end, self.epochs)\n\n    def on_epoch_begin(self, epoch, logs=None):\n        self.model.lambda_loss = self.lambda_values[epoch]\n        self.model.gamma_loss = self.gamma_values[epoch]","metadata":{"execution":{"iopub.status.busy":"2023-08-05T22:44:58.493395Z","iopub.execute_input":"2023-08-05T22:44:58.493774Z","iopub.status.idle":"2023-08-05T22:44:58.503979Z","shell.execute_reply.started":"2023-08-05T22:44:58.493748Z","shell.execute_reply":"2023-08-05T22:44:58.503214Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Objects for training <a class=\"anchor\" id=\"objects\"></a>","metadata":{"id":"52IPZaJLMRgE"}},{"cell_type":"markdown","source":"Create necessary objects.","metadata":{}},{"cell_type":"code","source":"EPOCHS = 43\nSTEPS_PER_EPOCH = 3000\nLAMBDA_START = 3\nLAMBDA_END = 1e-4\nGAMMA_START = 1e-4\nGAMMA_END = 0.999\nLAMBDA_ID = 1e-4\nEPOCH_INTERVAL_PLOT = 5\nEPOCH_INTERVAL_FID = None\nLR_START = 2e-4\nLR_END = 5e-6","metadata":{"execution":{"iopub.status.busy":"2023-08-05T22:44:58.504975Z","iopub.execute_input":"2023-08-05T22:44:58.505258Z","iopub.status.idle":"2023-08-05T22:44:58.521360Z","shell.execute_reply.started":"2023-08-05T22:44:58.505233Z","shell.execute_reply":"2023-08-05T22:44:58.520480Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with strategy.scope():\n    update_lr_cb = UpdateLearningRateCallback(EPOCHS, lr_start=LR_START, lr_end=LR_END)","metadata":{"execution":{"iopub.status.busy":"2023-08-05T22:44:58.522464Z","iopub.execute_input":"2023-08-05T22:44:58.522762Z","iopub.status.idle":"2023-08-05T22:44:58.534989Z","shell.execute_reply.started":"2023-08-05T22:44:58.522736Z","shell.execute_reply":"2023-08-05T22:44:58.534114Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with strategy.scope():\n    generator_g = create_generator()\n    generator_f = create_generator()\n    discriminator_x = create_discriminator(add_noise=True)\n    discriminator_y = create_discriminator(add_noise=True)","metadata":{"execution":{"iopub.status.busy":"2023-08-05T22:44:58.536121Z","iopub.execute_input":"2023-08-05T22:44:58.536449Z","iopub.status.idle":"2023-08-05T22:45:10.323791Z","shell.execute_reply.started":"2023-08-05T22:44:58.536420Z","shell.execute_reply":"2023-08-05T22:45:10.322681Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"generator_g.summary()","metadata":{"execution":{"iopub.status.busy":"2023-08-05T22:50:35.119930Z","iopub.execute_input":"2023-08-05T22:50:35.120982Z","iopub.status.idle":"2023-08-05T22:50:35.179222Z","shell.execute_reply.started":"2023-08-05T22:50:35.120947Z","shell.execute_reply":"2023-08-05T22:50:35.177962Z"},"_kg_hide-output":true,"collapsed":true,"jupyter":{"outputs_hidden":true},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"discriminator_x.summary()","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2023-08-05T22:53:10.358580Z","iopub.execute_input":"2023-08-05T22:53:10.359266Z","iopub.status.idle":"2023-08-05T22:53:10.389459Z","shell.execute_reply.started":"2023-08-05T22:53:10.359235Z","shell.execute_reply":"2023-08-05T22:53:10.388321Z"},"collapsed":true,"jupyter":{"outputs_hidden":true},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_pred_photo = next(iter(test_photo_ds.take(1)))\nplot_pred_photo = np.expand_dims(plot_pred_photo[0], axis=0)\nwith strategy.scope():\n    plot_pred_cb = PlotPredictionsCallback(\n        input_image=plot_pred_photo,\n        model_generator=generator_g,\n        epoch_interval=EPOCH_INTERVAL_PLOT)","metadata":{"id":"agbraWSxMRgG","execution":{"iopub.status.busy":"2023-08-05T22:45:10.324954Z","iopub.execute_input":"2023-08-05T22:45:10.325244Z","iopub.status.idle":"2023-08-05T22:45:10.407381Z","shell.execute_reply.started":"2023-08-05T22:45:10.325218Z","shell.execute_reply":"2023-08-05T22:45:10.406169Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fid_photo_ds = load_dataset(photo_files_tfrec).take(1024).batch(64).prefetch(AUTOTUNE)\nfid_monet_ds = load_dataset(monet_files_tfrec).batch(64).prefetch(AUTOTUNE)","metadata":{"id":"bIdXefNqMRgH","execution":{"iopub.status.busy":"2023-08-05T22:45:10.408727Z","iopub.execute_input":"2023-08-05T22:45:10.409057Z","iopub.status.idle":"2023-08-05T22:45:10.464283Z","shell.execute_reply.started":"2023-08-05T22:45:10.409028Z","shell.execute_reply":"2023-08-05T22:45:10.463294Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with strategy.scope():\n    fid_model = create_fid_inception_model()\n    fid_calc = FIDCalculator(\n        images_x_ds=fid_photo_ds,\n        images_y_ds=fid_monet_ds,\n        model_generator=generator_g,\n        fid_model_base=fid_model)\n    fid_calc.init_stat_x()\n    fid_cb = FIDCallback(fid_calculator=fid_calc, epoch_interval=EPOCH_INTERVAL_FID)","metadata":{"id":"ueutu8udMRgH","outputId":"7f0cfebb-dd63-4c6e-ad92-df40d7e61a21","_kg_hide-output":true,"execution":{"iopub.status.busy":"2023-08-05T22:45:10.465341Z","iopub.execute_input":"2023-08-05T22:45:10.465621Z","iopub.status.idle":"2023-08-05T22:45:56.514751Z","shell.execute_reply.started":"2023-08-05T22:45:10.465596Z","shell.execute_reply":"2023-08-05T22:45:56.513472Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with strategy.scope():\n    update_weights_cb = UpdateLossWeightsCallback(\n        EPOCHS,\n        lambda_start=LAMBDA_START,\n        lambda_end=LAMBDA_END,\n        gamma_start=GAMMA_START,\n        gamma_end=GAMMA_END)","metadata":{"execution":{"iopub.status.busy":"2023-08-05T22:45:56.516044Z","iopub.execute_input":"2023-08-05T22:45:56.516380Z","iopub.status.idle":"2023-08-05T22:45:56.522757Z","shell.execute_reply.started":"2023-08-05T22:45:56.516351Z","shell.execute_reply":"2023-08-05T22:45:56.521764Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training <a class=\"anchor\" id=\"training\"></a>","metadata":{"id":"QZ42eBNmMRgH"}},{"cell_type":"code","source":"with strategy.scope():\n    model_cycleGAN = TrainingGAN(\n        generator_g=generator_g,\n        generator_f=generator_f,\n        discriminator_x=discriminator_x,\n        discriminator_y=discriminator_y,\n        lambda_loss=LAMBDA_START,\n        lambda_id_loss=LAMBDA_ID,\n        gamma_loss=GAMMA_START,\n        diffaug_fn=diffaug_fn\n    )\n    generator_g_optimizer = Adam(learning_rate=LR_START, beta_1=0.5)\n    discriminator_x_optimizer = Adam(learning_rate=LR_START, beta_1=0.5)\n    generator_f_optimizer = Adam(learning_rate=LR_START, beta_1=0.5)\n    discriminator_y_optimizer = Adam(learning_rate=LR_START, beta_1=0.5)\n\n    model_cycleGAN.compile(\n        generator_g_optimizer=generator_g_optimizer,\n        discriminator_x_optimizer=discriminator_x_optimizer,\n        generator_f_optimizer=generator_f_optimizer,\n        discriminator_y_optimizer=discriminator_y_optimizer,\n    )","metadata":{"execution":{"iopub.status.busy":"2023-08-05T22:45:56.523840Z","iopub.execute_input":"2023-08-05T22:45:56.524121Z","iopub.status.idle":"2023-08-05T22:45:56.831590Z","shell.execute_reply.started":"2023-08-05T22:45:56.524096Z","shell.execute_reply":"2023-08-05T22:45:56.830178Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\nmodel_cycleGAN.fit(\n    train_ds,\n    epochs=EPOCHS,\n    callbacks=[plot_pred_cb, fid_cb, update_weights_cb, update_lr_cb],\n    steps_per_epoch=STEPS_PER_EPOCH\n)","metadata":{"id":"OIAA6r2wMRgI","outputId":"44217e46-0647-47a1-bd31-88f4f59bc943","execution":{"iopub.status.busy":"2023-08-05T22:45:56.832874Z","iopub.execute_input":"2023-08-05T22:45:56.833165Z","iopub.status.idle":"2023-08-05T22:50:06.711395Z","shell.execute_reply.started":"2023-08-05T22:45:56.833140Z","shell.execute_reply":"2023-08-05T22:50:06.710101Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Plot predictions <a class=\"anchor\" id=\"plotpredictions\"></a>","metadata":{"id":"2W7-r9O4MRgM"}},{"cell_type":"code","source":"def plot_predictions(model, ds):\n    ds_iter = iter(ds)\n    for n_sample in range(8):\n        example_sample = next(ds_iter)\n        generated_sample = model.predict(example_sample, verbose=0)\n        f = plt.figure(figsize=(13, 13))\n        plt.subplot(121)\n        plt.title('Input image')\n        plt.imshow(example_sample[0] * 0.5 + 0.5)\n        plt.axis('off')\n        plt.subplot(122)\n        plt.title('Generated image')\n        plt.imshow(generated_sample[0] * 0.5 + 0.5)\n        plt.axis('off')\n        plt.show()","metadata":{"id":"ZvdbAD2NMRgM","execution":{"iopub.status.busy":"2023-08-05T22:50:07.149899Z","iopub.execute_input":"2023-08-05T22:50:07.150260Z","iopub.status.idle":"2023-08-05T22:50:07.164238Z","shell.execute_reply.started":"2023-08-05T22:50:07.150228Z","shell.execute_reply":"2023-08-05T22:50:07.162923Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_predictions(generator_g, test_photo_ds.shuffle(1000))","metadata":{"id":"Y-g-HxD7MRgM","outputId":"9f597efa-33ad-4cae-e59a-085841c3b07a","execution":{"iopub.status.busy":"2023-08-05T22:50:07.165375Z","iopub.execute_input":"2023-08-05T22:50:07.165680Z","iopub.status.idle":"2023-08-05T22:50:28.548387Z","shell.execute_reply.started":"2023-08-05T22:50:07.165653Z","shell.execute_reply":"2023-08-05T22:50:28.546492Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Submission <a class=\"anchor\" id=\"submission\"></a>","metadata":{"id":"MbPPPH-eMRgN"}},{"cell_type":"code","source":"def generate_and_save(generator, ds):\n    count = 0\n    if os.path.exists('images.zip'):\n        os.remove('images.zip')\n    with zipfile.ZipFile('images.zip', 'w') as zipf:\n        for images_batch in ds:\n            predictions = generator(images_batch, training=False)\n            for pred in predictions:\n                count += 1\n                generated_image = array_to_img(pred)\n                generated_image_bytes = BytesIO()\n                generated_image.save(generated_image_bytes, format='JPEG')\n                generated_image_bytes.seek(0)\n                zipf.writestr(f'generated_image_{count}.jpg', generated_image_bytes.getvalue())\n                if count % 1024 == 0:\n                    print(f'Archived images: {count}')","metadata":{"id":"0pChjahqMRgO","execution":{"iopub.status.busy":"2023-08-05T22:50:28.551084Z","iopub.status.idle":"2023-08-05T22:50:28.551529Z","shell.execute_reply.started":"2023-08-05T22:50:28.551291Z","shell.execute_reply":"2023-08-05T22:50:28.551311Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\ngenerate_and_save(generator_g, test_photo_ds)","metadata":{"id":"bPtTNii4MRgO","execution":{"iopub.status.busy":"2023-08-05T22:50:28.552979Z","iopub.status.idle":"2023-08-05T22:50:28.553379Z","shell.execute_reply.started":"2023-08-05T22:50:28.553189Z","shell.execute_reply":"2023-08-05T22:50:28.553207Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if IN_COLAB:\n    from google.colab import drive\n    drive.mount('/content/gdrive')\n    !cp images.zip /content/gdrive/MyDrive/images.zip","metadata":{"_kg_hide-input":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Conclusion <a class=\"anchor\" id=\"conclusion\"></a>","metadata":{"id":"1LATMVMhMRgO"}},{"cell_type":"markdown","source":"Custom implementation using TensorFlow showed relatively good performance for modified CycleGAN. Among all the parameters, tuning the following ones has helped to improve the performance:\n* Reducing LAMBDA_START from 10 to 3\n* Add GaussianNoise layers after Conv2D and before InstanceNormalization layers, and to each downsample layer\n* Setting stddev from 0.1 to 0.2 for the GaussianNoise layer in the discriminator\n* Disable shuffling and caching in the dataset\n\nThe paper also suggests many other ideas for future work, such as how to tune the parameters, using a single discriminator for both directions, and so on.","metadata":{"id":"_PvfMDziMRgP"}},{"cell_type":"markdown","source":"# References <a class=\"anchor\" id=\"references\"></a>","metadata":{"id":"j1zJlB1KMRgP"}},{"cell_type":"markdown","source":"[Paper: Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks](https://arxiv.org/abs/1703.10593)\n\n[Paper: CycleGAN with Better Cycles](https://www.tongzhouwang.info/better_cycles/report.pdf)\n\n[Paper: Improved Techniques for Training GANs](https://arxiv.org/abs/1606.03498)\n\n[Paper: Differentiable Augmentation for Data-Efficient GAN Training](https://arxiv.org/abs/2006.10738)\n\n[GitHub: Transforming the World Into Paintings with CycleGAN](https://github.com/sebtheiler/tutorials/tree/main/cyclegan)\n\n[Article: pix2pix: Image-to-image translation with a conditional GAN](https://www.tensorflow.org/tutorials/generative/pix2pix)\n\n[Notebook: Two-objective discriminator by UnfriendlyAI](https://www.kaggle.com/code/unfriendlyai/two-objective-discriminator)","metadata":{"id":"2SX6LFWPMRgP"}}]}