{"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":"Resource:\n\n* [CS109B Data Science 2(Harvard) Generative Adversarial Networks (GANs) Vincent Casser, Pavlos Protopapas](https://harvard-iacs.github.io/2019-CS109B/a-sections/a-section8/presentation/cs109b_asec8_slides_gan.pdf)\n* [How to Develop a CycleGAN for Image-to-Image Translation with Keras](https://machinelearningmastery.com/cyclegan-tutorial-with-keras/)\n* [Image2Art Translation Using CycleGAN](https://sahiltinky94.medium.com/image2art-translation-using-cyclegan-e1bc096b7315)\n* [Transforming the World Into Paintings with CycleGAN](https://medium.com/analytics-vidhya/transforming-the-world-into-paintings-with-cyclegan-6748c0b85632)\n* [Tensorflow CycleGAN](https://www.tensorflow.org/tutorials/generative/cyclegan)\n* [Keras CycleGAN](https://keras.io/examples/generative/cyclegan/)  \n* [CycleGAN-Tensorflow-2 ResNet](https://github.com/LynnHo/CycleGAN-Tensorflow-2) \n* [CycleGAN-Tensorflow-2 U-net](https://github.com/LynnHo/AttGAN-Tensorflow) \n","metadata":{}},{"cell_type":"code","source":"import os, random, time, math, datetime, glob, sys, warnings, tqdm, functools, gc\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom IPython.display import clear_output\nimport tensorflow as tf\nfrom tensorflow.keras import backend as K\nfrom kaggle_datasets import KaggleDatasets\nfrom tensorflow.keras.optimizers import Adam, SGD, RMSprop\nfrom tensorflow.keras.models import Model, Sequential\nfrom tensorflow.keras.layers import Conv2D, Conv2DTranspose, LeakyReLU, Concatenate, Input,  ReLU, Layer, Dropout, ZeroPadding2D\nimport tensorflow_addons as tfa\nimport tensorflow.keras as keras\nAUTO = tf.data.experimental.AUTOTUNE \n\nwarnings.filterwarnings(\"ignore\")","metadata":{"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2021-10-05T11:31:52.660418Z","iopub.execute_input":"2021-10-05T11:31:52.660817Z","iopub.status.idle":"2021-10-05T11:31:57.302490Z","shell.execute_reply.started":"2021-10-05T11:31:52.660734Z","shell.execute_reply":"2021-10-05T11:31:57.301507Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"try:\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\nexcept ValueError:\n    tpu = None\n    gpus = tf.config.experimental.list_logical_devices(\"GPU\")\nif tpu:\n    tf.config.experimental_connect_to_cluster(tpu)\n    tf.tpu.experimental.initialize_tpu_system(tpu)\n    strategy = tf.distribute.experimental.TPUStrategy(tpu)\n    print('Running on TPU ', tpu.master())  \nelif len(gpus) > 0:\n    strategy = tf.distribute.MirroredStrategy(gpus)\n    print('Running on ', len(gpus), ' GPU(s) ')\nelse:\n    strategy = tf.distribute.get_strategy()\n    print('Running on CPU')\nprint(\"Number of accelerators: \", strategy.num_replicas_in_sync)","metadata":{"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2021-10-05T11:31:57.304253Z","iopub.execute_input":"2021-10-05T11:31:57.304624Z","iopub.status.idle":"2021-10-05T11:31:58.995241Z","shell.execute_reply.started":"2021-10-05T11:31:57.304584Z","shell.execute_reply":"2021-10-05T11:31:58.994353Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# GCS_DS_PATH = KaggleDatasets().get_gcs_path('gan-getting-started') \nGCS_DS_PATH = '../input/gan-getting-started'\nMONET_FILENAMES = tf.io.gfile.glob(GCS_DS_PATH+'/monet_jpg/*.jpg') # monet_tfrec monet_jpg\nPHOTO_FILENAMES = tf.io.gfile.glob(GCS_DS_PATH+'/photo_jpg/*.jpg') # photo_tfrec photo_jpg\n\nrandom.shuffle(MONET_FILENAMES)\nrandom.shuffle(PHOTO_FILENAMES)\n\nPHOTO_FILENAMES_TFREC = tf.io.gfile.glob(GCS_DS_PATH+'/photo_tfrec/*.tfrec')\nMONET_FILENAMES_TFREC = tf.io.gfile.glob(GCS_DS_PATH+'/monet_tfrec/*.tfrec')\n\nrandom.shuffle(PHOTO_FILENAMES_TFREC)\nrandom.shuffle(MONET_FILENAMES_TFREC)\n\nBATCH_SIZE = 1 * strategy.num_replicas_in_sync\nSTEPS_PER_EPOCH = 400\nEPOCHS = 200\n\nBUFFER_SIZE = 300\nIMAGE_HEIGHT = 256\nIMAGE_WIDTH  = 256\n\nDATA_SIZE = BATCH_SIZE*STEPS_PER_EPOCH\nIMAGE_SIZE = [IMAGE_HEIGHT, IMAGE_WIDTH]\nSAVE_FREQ = STEPS_PER_EPOCH * int(EPOCHS*0.2)\nif K.image_data_format() == 'channels_first':\n    SHAPE = (3,*IMAGE_SIZE)\nelse:\n    SHAPE = (*IMAGE_SIZE, 3)\n\nDIM = 64  # 64\nSIZE = 4 # 4\nSTART_LEARNING_RATE=3e-5\nSTART_LEARNING_RATE_DISC=1e-5 \nN_DOWNSAMPLINGS = 2 # 2\nN_DOWNSAMPLINGS_DESC = 3 # 3\nN_BLOCKS = 9 # 9","metadata":{"execution":{"iopub.status.busy":"2021-10-05T11:31:58.997109Z","iopub.execute_input":"2021-10-05T11:31:58.997657Z","iopub.status.idle":"2021-10-05T11:32:03.246432Z","shell.execute_reply.started":"2021-10-05T11:31:58.997618Z","shell.execute_reply":"2021-10-05T11:32:03.245389Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def decode_image(image):\n    image = tf.image.decode_jpeg(image, channels=3)\n    image = (tf.cast(image, tf.float32) / 127.5) - 1\n    return tf.reshape(image, SHAPE)\n\n@tf.function\ndef _map_fn(img):  # preprocessing\n    img = tf.image.resize(img, [286, 286])  # or img = tf.image.resize(img, [load_size, load_size]); img = tl.center_crop(img, crop_size)\n    img = tf.clip_by_value(img, 0, 255) / 255.0  # or img = tl.minmax_norm(img)\n    img = img * 2 - 1\n    return img\n\n@tf.function\ndef augmentation(image):\n    aug_size = IMAGE_WIDTH+int(IMAGE_WIDTH*0.15)\n    image = tf.image.resize(image, [aug_size, aug_size], method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)\n    image = tf.image.random_crop(image, size=[IMAGE_WIDTH, IMAGE_HEIGHT, 3])\n    image = tf.image.random_flip_up_down(image)\n    image = tf.image.random_flip_left_right(image)\n    return image\n\ndef read_tfrecord(example):\n    tfrecord_format = { \"image\": tf.io.FixedLenFeature([], tf.string) }\n    example = tf.io.parse_single_example(example, tfrecord_format)\n    return decode_image(example['image'])\n     \ndef load_dataset(filenames):\n    dataset = tf.data.TFRecordDataset(filenames)\n    return dataset.map(read_tfrecord, num_parallel_calls=AUTO)\n\n\nfast_photo_ds = load_dataset(PHOTO_FILENAMES_TFREC).shuffle(BUFFER_SIZE).batch(32*strategy.num_replicas_in_sync, drop_remainder=True).prefetch(AUTO)\nfid_photo_ds = load_dataset(PHOTO_FILENAMES_TFREC).shuffle(BUFFER_SIZE).take(1024).batch(32*strategy.num_replicas_in_sync, drop_remainder=True).prefetch(AUTO)\nfid_monet_ds = load_dataset(MONET_FILENAMES_TFREC).shuffle(BUFFER_SIZE).batch(32*strategy.num_replicas_in_sync, drop_remainder=True).prefetch(AUTO)\n\ndef get_gan_dataset():\n    monet_ds = load_dataset(MONET_FILENAMES_TFREC)\n    photo_ds = load_dataset(PHOTO_FILENAMES_TFREC)\n    monet_ds = monet_ds.shuffle(BUFFER_SIZE).batch(BATCH_SIZE, drop_remainder=True).cache().shuffle(BUFFER_SIZE).repeat().prefetch(AUTO)\n    photo_ds = photo_ds.shuffle(BUFFER_SIZE).batch(BATCH_SIZE, drop_remainder=True).cache().shuffle(BUFFER_SIZE).repeat().prefetch(AUTO)\n    return tf.data.Dataset.zip((monet_ds, photo_ds))","metadata":{"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2021-10-05T11:32:03.248069Z","iopub.execute_input":"2021-10-05T11:32:03.248426Z","iopub.status.idle":"2021-10-05T11:32:03.437569Z","shell.execute_reply.started":"2021-10-05T11:32:03.248395Z","shell.execute_reply":"2021-10-05T11:32:03.436731Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import multiprocessing\ndef batch_dataset(dataset,\n                  batch_size,\n                  drop_remainder=True,\n                  n_prefetch_batch=1,\n                  filter_fn=None,\n                  map_fn=None,\n                  n_map_threads=None,\n                  filter_after_map=False,\n                  shuffle=True,\n                  shuffle_buffer_size=None,\n                  repeat=AUTO):\n    # set defaults\n    if n_map_threads is None:\n        n_map_threads = multiprocessing.cpu_count()\n    if shuffle and shuffle_buffer_size is None:\n        shuffle_buffer_size = max(batch_size * 128, 2048)  # set the minimum buffer size as 2048\n\n    # [*] it is efficient to conduct `shuffle` before `map`/`filter` because `map`/`filter` is sometimes costly\n    if shuffle:\n        dataset = dataset.shuffle(shuffle_buffer_size)\n\n    if not filter_after_map:\n        if filter_fn:\n            dataset = dataset.filter(filter_fn)\n\n        if map_fn:\n            dataset = dataset.map(map_fn, num_parallel_calls=n_map_threads)\n\n    else:  # [*] this is slower\n        if map_fn:\n            dataset = dataset.map(map_fn, num_parallel_calls=n_map_threads)\n\n        if filter_fn:\n            dataset = dataset.filter(filter_fn)\n\n    dataset = dataset.batch(batch_size, drop_remainder=drop_remainder)\n\n    dataset = dataset.repeat().prefetch(n_prefetch_batch)\n\n    return dataset\n\n\ndef memory_data_batch_dataset(memory_data,\n                              batch_size,\n                              drop_remainder=True,\n                              n_prefetch_batch=1,\n                              filter_fn=None,\n                              map_fn=None,\n                              n_map_threads=None,\n                              filter_after_map=False,\n                              shuffle=True,\n                              shuffle_buffer_size=None,\n                              repeat=AUTO):\n    \"\"\"Batch dataset of memory data.\n    Parameters\n    ----------\n    memory_data : nested structure of tensors/ndarrays/lists\n    \"\"\"\n    dataset = tf.data.Dataset.from_tensor_slices(memory_data)\n    dataset = batch_dataset(dataset,\n                            batch_size,\n                            drop_remainder=drop_remainder,\n                            n_prefetch_batch=n_prefetch_batch,\n                            filter_fn=filter_fn,\n                            map_fn=map_fn,\n                            n_map_threads=n_map_threads,\n                            filter_after_map=filter_after_map,\n                            shuffle=shuffle,\n                            shuffle_buffer_size=shuffle_buffer_size,\n                            repeat=repeat)\n    return dataset\n\n\ndef disk_image_batch_dataset(img_paths,\n                             batch_size,\n                             labels=None,\n                             drop_remainder=True,\n                             n_prefetch_batch=1,\n                             filter_fn=None,\n                             map_fn=None,\n                             n_map_threads=None,\n                             filter_after_map=False,\n                             shuffle=True,\n                             shuffle_buffer_size=None,\n                             repeat=AUTO):\n    \"\"\"Batch dataset of disk image for PNG and JPEG.\n    Parameters\n    ----------\n    img_paths : 1d-tensor/ndarray/list of str\n    labels : nested structure of tensors/ndarrays/lists\n    \"\"\"\n    if labels is None:\n        memory_data = img_paths\n    else:\n        memory_data = (img_paths, labels)\n\n    def parse_fn(path):\n        img = tf.io.read_file(path)\n        img = tf.image.decode_jpeg(img, 3)  # fix channels to 3\n        img = tf.cast(img, tf.float32)\n        return (img,)\n\n    if map_fn:  # fuse `map_fn` and `parse_fn`\n        def map_fn_(*args):\n            return map_fn(*parse_fn(*args))\n    else:\n        map_fn_ = parse_fn\n\n    dataset = memory_data_batch_dataset(memory_data,\n                                        batch_size,\n                                        drop_remainder=drop_remainder,\n                                        n_prefetch_batch=n_prefetch_batch,\n                                        filter_fn=filter_fn,\n                                        map_fn=map_fn_,\n                                        n_map_threads=n_map_threads,\n                                        filter_after_map=filter_after_map,\n                                        shuffle=shuffle,\n                                        shuffle_buffer_size=shuffle_buffer_size,\n                                        repeat=repeat)\n\n    return dataset\n\ndef make_dataset(img_paths, batch_size, load_size, crop_size, training, drop_remainder=True, shuffle=True, repeat=AUTO):\n    if training:\n        @tf.function\n        def _map_fn(img):  # preprocessing\n            img = tf.image.random_flip_left_right(img)\n            img = tf.image.resize(img, [load_size, load_size])\n            img = tf.image.random_crop(img, [crop_size, crop_size, tf.shape(img)[-1]])\n            img = tf.clip_by_value(img, 0, 255) / 255.0  # or img = tl.minmax_norm(img)\n            img = img * 2 - 1\n            return img\n    else:\n        @tf.function\n        def _map_fn(img):  # preprocessing\n            img = tf.image.resize(img, [crop_size, crop_size])  # or img = tf.image.resize(img, [load_size, load_size]); img = tl.center_crop(img, crop_size)\n            img = tf.clip_by_value(img, 0, 255) / 255.0  # or img = tl.minmax_norm(img)\n            img = img * 2 - 1\n            return img\n\n    return disk_image_batch_dataset(img_paths,\n                                       batch_size,\n                                       drop_remainder=drop_remainder,\n                                       map_fn=_map_fn,\n                                       shuffle=shuffle,\n                                       repeat=repeat)\n\n\ndef make_zip_dataset(A_img_paths=MONET_FILENAMES, B_img_paths=PHOTO_FILENAMES, batch_size=BATCH_SIZE, load_size=IMAGE_WIDTH+int(IMAGE_WIDTH*0.15), crop_size=IMAGE_WIDTH, training=True, shuffle=True, repeat=AUTO):\n    # zip two datasets aligned by the longer one\n    if repeat:\n        A_repeat = B_repeat = None  # cycle both\n    else:\n        if len(A_img_paths) >= len(B_img_paths):\n            A_repeat = 1\n            B_repeat = None  # cycle the shorter one\n        else:\n            A_repeat = None  # cycle the shorter one\n            B_repeat = 1\n\n    A_dataset = make_dataset(A_img_paths, batch_size, load_size, crop_size, training, drop_remainder=True, shuffle=shuffle, repeat=A_repeat)\n    B_dataset = make_dataset(B_img_paths, batch_size, load_size, crop_size, training, drop_remainder=True, shuffle=shuffle, repeat=B_repeat)\n\n    A_B_dataset = tf.data.Dataset.zip((A_dataset, B_dataset))\n    len_dataset = max(len(A_img_paths), len(B_img_paths)) // batch_size\n\n    return A_B_dataset, len_dataset\n\n\nclass ItemPool:\n    def __init__(self, pool_size=50):\n        self.pool_size = pool_size\n        self.items = []\n\n    def __call__(self, in_items):\n        # `in_items` should be a batch tensor\n\n        if self.pool_size == 0:\n            return in_items\n\n        out_items = []\n        for in_item in in_items:\n            if len(self.items) < self.pool_size:\n                self.items.append(in_item)\n                out_items.append(in_item)\n            else:\n                if np.random.rand() > 0.5:\n                    idx = np.random.randint(0, len(self.items))\n                    out_item, self.items[idx] = self.items[idx], in_item\n                    out_items.append(out_item)\n                else:\n                    out_items.append(in_item)\n        return tf.stack(out_items, axis=0)","metadata":{"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2021-10-05T11:32:03.438833Z","iopub.execute_input":"2021-10-05T11:32:03.439155Z","iopub.status.idle":"2021-10-05T11:32:03.468597Z","shell.execute_reply.started":"2021-10-05T11:32:03.439121Z","shell.execute_reply":"2021-10-05T11:32:03.467815Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dataset, len_dataset = make_zip_dataset()\ntest_dataset=get_gan_dataset()","metadata":{"execution":{"iopub.status.busy":"2021-10-05T11:32:03.470137Z","iopub.execute_input":"2021-10-05T11:32:03.470551Z","iopub.status.idle":"2021-10-05T11:32:03.788994Z","shell.execute_reply.started":"2021-10-05T11:32:03.470503Z","shell.execute_reply":"2021-10-05T11:32:03.787927Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def plot_images(sample_monet1, sample_monet2):\n    plt.figure(figsize=(16, 816))\n    display_list = [sample_monet2[0], sample_monet1[0]]\n    title = ['photo', 'monet']\n\n    for i in range(2):\n        plt.subplot(1, 2, i+1)\n        plt.title(title[i])\n        plt.imshow(display_list[i] * 0.5 + 0.5)\n        plt.axis('off')\n    plt.show()\n    \nfor monet, photo in dataset.take(2):\n    plot_images(monet, photo)","metadata":{"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2021-10-05T11:32:03.790437Z","iopub.execute_input":"2021-10-05T11:32:03.790844Z","iopub.status.idle":"2021-10-05T11:32:05.435102Z","shell.execute_reply.started":"2021-10-05T11:32:03.790801Z","shell.execute_reply":"2021-10-05T11:32:05.434169Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Losses fn","metadata":{}},{"cell_type":"code","source":"def get_gan_losses_fn():\n    bce = tf.losses.BinaryCrossentropy(from_logits=True)\n\n    def d_loss_fn(r_logit, f_logit):\n        r_loss = bce(tf.ones_like(r_logit), r_logit)\n        f_loss = bce(tf.zeros_like(f_logit), f_logit)\n        return r_loss, f_loss\n\n    def g_loss_fn(f_logit):\n        f_loss = bce(tf.ones_like(f_logit), f_logit)\n        return f_loss\n\n    return d_loss_fn, g_loss_fn\n\n\ndef get_hinge_v1_losses_fn():\n    def d_loss_fn(r_logit, f_logit):\n        r_loss = tf.reduce_mean(tf.maximum(1 - r_logit, 0))\n        f_loss = tf.reduce_mean(tf.maximum(1 + f_logit, 0))\n        return r_loss, f_loss\n\n    def g_loss_fn(f_logit):\n        f_loss = tf.reduce_mean(tf.maximum(1 - f_logit, 0))\n        return f_loss\n\n    return d_loss_fn, g_loss_fn\n\n\ndef get_hinge_v2_losses_fn():\n    def d_loss_fn(r_logit, f_logit):\n        r_loss = tf.reduce_mean(tf.maximum(1 - r_logit, 0))\n        f_loss = tf.reduce_mean(tf.maximum(1 + f_logit, 0))\n        return r_loss, f_loss\n\n    def g_loss_fn(f_logit):\n        f_loss = tf.reduce_mean(- f_logit)\n        return f_loss\n\n    return d_loss_fn, g_loss_fn\n\n\ndef get_lsgan_losses_fn():\n    mse = tf.losses.MeanSquaredError()\n\n    def d_loss_fn(r_logit, f_logit):\n        r_loss = mse(tf.ones_like(r_logit), r_logit)\n        f_loss = mse(tf.zeros_like(f_logit), f_logit)\n        return r_loss, f_loss\n\n    def g_loss_fn(f_logit):\n        f_loss = mse(tf.ones_like(f_logit), f_logit)\n        return f_loss\n\n    return d_loss_fn, g_loss_fn\n\n\ndef get_wgan_losses_fn():\n    def d_loss_fn(r_logit, f_logit):\n        r_loss = - tf.reduce_mean(r_logit)\n        f_loss = tf.reduce_mean(f_logit)\n        return r_loss, f_loss\n\n    def g_loss_fn(f_logit):\n        f_loss = - tf.reduce_mean(f_logit)\n        return f_loss\n\n    return d_loss_fn, g_loss_fn\n\n\ndef get_adversarial_losses_fn(mode):\n    if mode == 'gan':\n        return get_gan_losses_fn()\n    elif mode == 'hinge_v1':\n        return get_hinge_v1_losses_fn()\n    elif mode == 'hinge_v2':\n        return get_hinge_v2_losses_fn()\n    elif mode == 'lsgan':\n        return get_lsgan_losses_fn()\n    elif mode == 'wgan':\n        return get_wgan_losses_fn()\n@tf.function\ndef minmax_norm(x, epsilon=1e-12):\n    x = tf.cast(x, tf.float32)\n    min_val = tf.reduce_min(x)\n    max_val = tf.reduce_max(x)\n    norm_x = (x - min_val) / tf.maximum((max_val - min_val), epsilon)\n    return norm_x\n\n\n@tf.function\ndef reshape(x, shape):\n    x = tf.convert_to_tensor(x)\n    shape = [x.shape[i] if shape[i] == 0 else shape[i] for i in range(len(shape))]  # TODO(Lynn): is it slow here?\n    shape = [tf.shape(x)[i] if shape[i] is None else shape[i] for i in range(len(shape))]\n    return tf.reshape(x, shape)\n\ndef gradient_penalty(f, real, fake, mode):\n    def _gradient_penalty(f, real, fake=None):\n        def _interpolate(a, b=None):\n            if b is None:   # interpolation in DRAGAN\n                beta = tf.random.uniform(shape=tf.shape(a), minval=0., maxval=1.)\n                b = a + 0.5 * tf.math.reduce_std(a) * beta\n            shape = [tf.shape(a)[0]] + [1] * (a.shape.ndims - 1)\n            alpha = tf.random.uniform(shape=shape, minval=0., maxval=1.)\n            inter = a + alpha * (b - a)\n            inter.set_shape(a.shape)\n            return inter\n\n        x = _interpolate(real, fake)\n        with tf.GradientTape() as t:\n            t.watch(x)\n            pred = f(x)\n        grad = t.gradient(pred, x)\n        norm = tf.norm(tf.reshape(grad, [tf.shape(grad)[0], -1]), axis=1)\n        gp = tf.reduce_mean((norm - 1.)**2)\n\n        return gp\n\n    if mode == 'none':\n        gp = tf.constant(0, dtype=real.dtype)\n    elif mode == 'dragan':\n        gp = _gradient_penalty(f, real)\n    elif mode == 'wgan-gp':\n        gp = _gradient_penalty(f, real, fake)\n\n    return gp","metadata":{"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2021-10-05T11:32:05.438748Z","iopub.execute_input":"2021-10-05T11:32:05.439082Z","iopub.status.idle":"2021-10-05T11:32:05.608483Z","shell.execute_reply.started":"2021-10-05T11:32:05.439049Z","shell.execute_reply":"2021-10-05T11:32:05.606998Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### make dataset @tf.function","metadata":{}},{"cell_type":"code","source":"# import multiprocessing\n# def batch_dataset(dataset,\n#                   batch_size,\n#                   drop_remainder=True,\n#                   n_prefetch_batch=1,\n#                   filter_fn=None,\n#                   map_fn=None,\n#                   n_map_threads=None,\n#                   filter_after_map=False,\n#                   shuffle=True,\n#                   shuffle_buffer_size=None,\n#                   repeat=None):\n#     # set defaults\n#     if n_map_threads is None:\n#         n_map_threads = multiprocessing.cpu_count()\n#     if shuffle and shuffle_buffer_size is None:\n#         shuffle_buffer_size = max(batch_size * 128, 2048)  # set the minimum buffer size as 2048\n\n#     # [*] it is efficient to conduct `shuffle` before `map`/`filter` because `map`/`filter` is sometimes costly\n#     if shuffle:\n#         dataset = dataset.shuffle(shuffle_buffer_size)\n\n#     if not filter_after_map:\n#         if filter_fn:\n#             dataset = dataset.filter(filter_fn)\n\n#         if map_fn:\n#             dataset = dataset.map(map_fn, num_parallel_calls=n_map_threads)\n\n#     else:  # [*] this is slower\n#         if map_fn:\n#             dataset = dataset.map(map_fn, num_parallel_calls=n_map_threads)\n\n#         if filter_fn:\n#             dataset = dataset.filter(filter_fn)\n\n#     dataset = dataset.batch(batch_size, drop_remainder=drop_remainder)\n\n#     dataset = dataset.repeat(repeat).prefetch(n_prefetch_batch)\n\n#     return dataset\n\n\n# def memory_data_batch_dataset(memory_data,\n#                               batch_size,\n#                               drop_remainder=True,\n#                               n_prefetch_batch=1,\n#                               filter_fn=None,\n#                               map_fn=None,\n#                               n_map_threads=None,\n#                               filter_after_map=False,\n#                               shuffle=True,\n#                               shuffle_buffer_size=None,\n#                               repeat=None):\n#     \"\"\"Batch dataset of memory data.\n#     Parameters\n#     ----------\n#     memory_data : nested structure of tensors/ndarrays/lists\n#     \"\"\"\n#     dataset = tf.data.Dataset.from_tensor_slices(memory_data)\n#     dataset = batch_dataset(dataset,\n#                             batch_size,\n#                             drop_remainder=drop_remainder,\n#                             n_prefetch_batch=n_prefetch_batch,\n#                             filter_fn=filter_fn,\n#                             map_fn=map_fn,\n#                             n_map_threads=n_map_threads,\n#                             filter_after_map=filter_after_map,\n#                             shuffle=shuffle,\n#                             shuffle_buffer_size=shuffle_buffer_size,\n#                             repeat=repeat)\n#     return dataset\n\n\n# def disk_image_batch_dataset(img_paths,\n#                              batch_size,\n#                              labels=None,\n#                              drop_remainder=True,\n#                              n_prefetch_batch=1,\n#                              filter_fn=None,\n#                              map_fn=None,\n#                              n_map_threads=None,\n#                              filter_after_map=False,\n#                              shuffle=True,\n#                              shuffle_buffer_size=None,\n#                              repeat=None):\n#     dataset = load_dataset(img_paths)\n#     dataset = dataset.cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE, drop_remainder=True).repeat().prefetch(AUTO)\n\n#     return dataset\n\n# def make_dataset(img_paths, batch_size, load_size, crop_size, training, drop_remainder=True, shuffle=True, repeat=1):\n#     if training:\n#         @tf.function\n#         def _map_fn(img):  # preprocessing\n#             img = tf.image.random_flip_left_right(img)\n#             img = tf.image.resize(img, [load_size, load_size])\n#             img = tf.image.random_crop(img, [crop_size, crop_size, tf.shape(img)[-1]])\n#             img = tf.clip_by_value(img, 0, 255) / 255.0  # or img = tl.minmax_norm(img)\n#             img = img * 2 - 1\n#             return img\n#     else:\n#         @tf.function\n#         def _map_fn(img):  # preprocessing\n#             img = tf.image.resize(img, [crop_size, crop_size])  # or img = tf.image.resize(img, [load_size, load_size]); img = tl.center_crop(img, crop_size)\n#             img = tf.clip_by_value(img, 0, 255) / 255.0  # or img = tl.minmax_norm(img)\n#             img = img * 2 - 1\n#             return img\n\n#     return disk_image_batch_dataset(img_paths,\n#                                        batch_size,\n#                                        drop_remainder=drop_remainder,\n#                                        map_fn=_map_fn,\n#                                        shuffle=shuffle,\n#                                        repeat=repeat)\n\n\n# def make_zip_dataset(A_img_paths, B_img_paths, batch_size, load_size, crop_size, training, shuffle=True, repeat=False):\n#     # zip two datasets aligned by the longer one\n#     if repeat:\n#         A_repeat = B_repeat = None  # cycle both\n#     else:\n#         if len(A_img_paths) >= len(B_img_paths):\n#             A_repeat = 1\n#             B_repeat = None  # cycle the shorter one\n#         else:\n#             A_repeat = None  # cycle the shorter one\n#             B_repeat = 1\n\n#     A_dataset = make_dataset(A_img_paths, batch_size, load_size, crop_size, training, drop_remainder=True, shuffle=shuffle, repeat=A_repeat)\n#     B_dataset = make_dataset(B_img_paths, batch_size, load_size, crop_size, training, drop_remainder=True, shuffle=shuffle, repeat=B_repeat)\n\n#     A_B_dataset = tf.data.Dataset.zip((A_dataset, B_dataset))\n#     len_dataset = max(len(A_img_paths), len(B_img_paths)) // batch_size\n\n#     return A_B_dataset, len_dataset\n\n\n# class ItemPool:\n#     def __init__(self, pool_size=50):\n#         self.pool_size = pool_size\n#         self.items = []\n\n#     def __call__(self, in_items):\n#         # `in_items` should be a batch tensor\n\n#         if self.pool_size == 0:\n#             return in_items\n\n#         out_items = []\n#         for in_item in in_items:\n#             if len(self.items) < self.pool_size:\n#                 self.items.append(in_item)\n#                 out_items.append(in_item)\n#             else:\n#                 if np.random.rand() > 0.5:\n#                     idx = np.random.randint(0, len(self.items))\n#                     out_item, self.items[idx] = self.items[idx], in_item\n#                     out_items.append(out_item)\n#                 else:\n#                     out_items.append(in_item)\n#         return tf.stack(out_items, axis=0)","metadata":{"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2021-10-05T11:32:05.610659Z","iopub.execute_input":"2021-10-05T11:32:05.611253Z","iopub.status.idle":"2021-10-05T11:32:05.620991Z","shell.execute_reply.started":"2021-10-05T11:32:05.611196Z","shell.execute_reply":"2021-10-05T11:32:05.620131Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### ResnetGenerator ConvDiscriminator LinearDecay","metadata":{}},{"cell_type":"code","source":"def _get_norm_layer(norm):\n    if norm == 'none':\n        return lambda: lambda x: x\n    elif norm == 'batch_norm':\n        return keras.layers.BatchNormalization\n    elif norm == 'instance_norm':\n        return tfa.layers.InstanceNormalization\n    elif norm == 'layer_norm':\n        return keras.layers.LayerNormalization\n\ndef ResnetGenerator(input_shape=SHAPE,\n                    output_channels=3,\n                    dim=DIM,\n                    n_downsamplings=N_DOWNSAMPLINGS,\n                    n_blocks=N_BLOCKS, # 9\n                    norm='instance_norm'):\n    Norm = _get_norm_layer(norm)\n\n    def _residual_block(x):\n        dim = x.shape[-1]\n        h = x\n\n        h = tf.pad(h, [[0, 0], [1, 1], [1, 1], [0, 0]], mode='REFLECT')\n        h = keras.layers.Conv2D(dim, 3, padding='valid', use_bias=False)(h)\n        h = Norm()(h)\n        h = tf.nn.relu(h)\n\n        h = tf.pad(h, [[0, 0], [1, 1], [1, 1], [0, 0]], mode='REFLECT')\n        h = keras.layers.Conv2D(dim, 3, padding='valid', use_bias=False)(h)\n        h = Norm()(h)\n\n        return keras.layers.add([x, h])\n\n    # 0\n    h = inputs = keras.Input(shape=input_shape)\n\n    # 1\n    h = tf.pad(h, [[0, 0], [3, 3], [3, 3], [0, 0]], mode='REFLECT')\n    h = keras.layers.Conv2D(dim, 7, padding='valid', use_bias=False)(h)\n    h = Norm()(h)\n    h = tf.nn.relu(h)\n\n    # 2\n    for _ in range(n_downsamplings):\n        dim *= 2\n        h = keras.layers.Conv2D(dim, 3, strides=2, padding='same', use_bias=False)(h)\n        h = Norm()(h)\n        h = tf.nn.relu(h)\n\n    # 3\n    for _ in range(n_blocks):\n        h = _residual_block(h)\n\n    # 4\n    for _ in range(n_downsamplings):\n        dim //= 2\n        h = keras.layers.Conv2DTranspose(dim, 3, strides=2, padding='same', use_bias=False)(h)\n        h = Norm()(h)\n        h = tf.nn.relu(h)\n\n    # 5\n    h = tf.pad(h, [[0, 0], [3, 3], [3, 3], [0, 0]], mode='REFLECT')\n    h = keras.layers.Conv2D(output_channels, 7, padding='valid')(h)\n    h = tf.tanh(h)\n\n    return keras.Model(inputs=inputs, outputs=h)\n\ndef ConvDiscriminator(input_shape=SHAPE,\n                      dim=DIM,\n                      n_downsamplings=N_DOWNSAMPLINGS_DESC,\n                      norm='instance_norm'):\n    dim_ = dim\n    Norm = _get_norm_layer(norm)\n\n    # 0\n    h = inputs = keras.Input(shape=input_shape)\n\n    # 1\n    h = keras.layers.Conv2D(dim, 4, strides=2, padding='same')(h)\n    h = tf.nn.leaky_relu(h, alpha=0.2)\n\n    for _ in range(n_downsamplings - 1):\n        dim = min(dim * 2, dim_ * 8)\n        h = keras.layers.Conv2D(dim, 4, strides=2, padding='same', use_bias=False)(h)\n        h = Norm()(h)\n        h = tf.nn.leaky_relu(h, alpha=0.2)\n\n    # 2\n    dim = min(dim * 2, dim_ * 8)\n    h = keras.layers.Conv2D(dim, 4, strides=1, padding='same', use_bias=False)(h)\n    h = Norm()(h)\n    h = tf.nn.leaky_relu(h, alpha=0.2)\n\n    # 3\n    h = keras.layers.Conv2D(1, 4, strides=1, padding='same')(h)\n\n    return keras.Model(inputs=inputs, outputs=h)\n\nclass LinearDecay(keras.optimizers.schedules.LearningRateSchedule):\n    # if `step` < `step_decay`: use fixed learning rate\n    # else: linearly decay the learning rate to zero\n\n    def __init__(self, initial_learning_rate, total_steps, step_decay):\n        super(LinearDecay, self).__init__()\n        self._initial_learning_rate = initial_learning_rate\n        self._steps = total_steps\n        self._step_decay = step_decay\n        self.current_learning_rate = tf.Variable(initial_value=initial_learning_rate, trainable=False, dtype=tf.float32)\n\n    def __call__(self, step):\n        self.current_learning_rate.assign(tf.cond(\n            step >= self._step_decay,\n            true_fn=lambda: self._initial_learning_rate * (1 - 1 / (self._steps - self._step_decay) * (step - self._step_decay)),\n            false_fn=lambda: self._initial_learning_rate\n        ))\n        return self.current_learning_rate","metadata":{"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2021-10-05T11:32:05.622352Z","iopub.execute_input":"2021-10-05T11:32:05.622792Z","iopub.status.idle":"2021-10-05T11:32:05.649696Z","shell.execute_reply.started":"2021-10-05T11:32:05.622755Z","shell.execute_reply":"2021-10-05T11:32:05.648765Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Plot lr","metadata":{}},{"cell_type":"code","source":"# Matplotlib config\nplt.ioff()\nplt.rc('image', cmap='gray_r')\nplt.rc('grid', linewidth=1)\nplt.rc('xtick', top=False, bottom=False, labelsize='large')\nplt.rc('ytick', left=False, right=False, labelsize='large')\nplt.rc('axes', facecolor='F8F8F8', titlesize=\"large\", edgecolor='white')\nplt.rc('text', color='a8151a')\nplt.rc('figure', facecolor='F0F0F0', figsize=(16,9))\n# Matplotlib fonts\nMATPLOTLIB_FONT_DIR = os.path.join(os.path.dirname(plt.__file__), \"mpl-data/fonts/ttf\")\n\ndef plot_learning_rate(lr_func, epochs):\n    xx = np.arange(epochs+1, dtype=np.float)\n    y = [lr_decay(x) for x in xx]\n    fig, ax = plt.subplots(figsize=(9, 6))\n    ax.set_xlabel('epochs')\n    ax.set_title('Learning rate\\ndecays from {:0.3g} to {:0.3g}'.format(y[0], y[-2]))\n    ax.minorticks_on()\n    ax.grid(True, which='major', axis='both', linestyle='-', linewidth=1)\n    ax.grid(True, which='minor', axis='both', linestyle=':', linewidth=0.5)\n    ax.step(xx,y, linewidth=3, where='post')\n    display(fig)\n\ndef display_epoch_predict(photo2monet, photo, monet, monet2photo):\n    plt.figure(figsize=(16, 816))\n    display_list = [photo[0], photo2monet[0]]\n    title = [f'photo', f'photo2monet']\n    for i in range(2):\n        plt.subplot(1, 2, i+1)\n        plt.title(title[i])\n        plt.imshow(display_list[i] * 0.5 + 0.5)\n        plt.axis('off')\n    plt.show()\n    \n    plt.figure(figsize=(16, 816))\n    display_list = [monet[0], monet2photo[0]]\n    title = ['monet', 'monet2photo']\n    for i in range(2):\n        plt.subplot(1, 2, i+1)\n        plt.title(title[i])\n        plt.imshow(display_list[i] * 0.5 + 0.5)\n        plt.axis('off')\n    plt.show()\n         \nclass PlotTraining(tf.keras.callbacks.Callback):\n    def __init__(self, sample_rate=10, zoom=16):\n        self.sample_rate = sample_rate\n        self.step = 0\n        self.zoom = zoom\n        self.steps_per_epoch = STEPS_PER_EPOCH*BATCH_SIZE\n        \n    def on_train_begin(self, logs={}):\n        self.batch_history = {}\n        self.batch_step = []\n        self.epoch_history = {}\n        self.epoch_step = []\n        self.fig, self.axes = plt.subplots(1, 2, figsize=(16, 7))\n        self.fig.subplots_adjust(wspace=0.12, hspace=0.12)\n        plt.ioff()\n       \n    def on_batch_end(self, batch, logs={}):\n        if (batch % self.sample_rate) == 0:\n            self.batch_step.append(self.step)\n            for k,v in logs.items():\n                # do not log \"batch\" and \"size\" metrics that do not change\n                # do not log training accuracy \"acc\"\n                if k=='batch' or k=='size' or k.endswith('gen_loss'): # or k=='acc':\n                    continue\n                self.batch_history.setdefault(k, []).append(v)\n#                 self.batch_history.setdefault(k, []).append(v[-1][-1][-1])\n                self.step += 1\n \n    def on_epoch_end(self, epoch, logs={}):\n        plt.close(self.fig)\n        self.axes[0].cla()\n        self.axes[1].cla()\n    \n        self.axes[0].set_ylim(0, 2)\n        self.axes[1].set_ylim(0, 10)\n\n        self.epoch_step.append(self.step)\n     \n        for k,v in logs.items():\n          # only log validation metrics\n            if k.endswith('disc_loss'):\n                  continue\n            self.epoch_history.setdefault(k, []).append(v)\n#             self.epoch_history.setdefault(k, []).append(v[-1][-1][-1])\n        \n        for k,v in self.batch_history.items():\n            self.axes[0 if k.endswith('disc_loss') else 1].plot(np.array(self.batch_step) / self.steps_per_epoch, v, label=k)\n\n        for k,v in self.epoch_history.items():\n            self.axes[0 if k.endswith('disc_loss') else 1].plot(np.array(self.epoch_step) / self.steps_per_epoch, v, label=k, linewidth=3)\n\n        self.axes[0].legend()\n        self.axes[1].legend()\n        self.axes[0].set_xlabel('epochs')\n        self.axes[1].set_xlabel('epochs')\n        self.axes[0].minorticks_on()\n        self.axes[0].grid(True, which='major', axis='both', linestyle='-', linewidth=1)\n        self.axes[0].grid(True, which='minor', axis='both', linestyle=':', linewidth=0.5)\n        self.axes[1].minorticks_on()\n        self.axes[1].grid(True, which='major', axis='both', linestyle='-', linewidth=1)\n        self.axes[1].grid(True, which='minor', axis='both', linestyle=':', linewidth=0.5)\n        display(self.fig)","metadata":{"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2021-10-05T11:32:05.651157Z","iopub.execute_input":"2021-10-05T11:32:05.651605Z","iopub.status.idle":"2021-10-05T11:32:05.679396Z","shell.execute_reply.started":"2021-10-05T11:32:05.651546Z","shell.execute_reply":"2021-10-05T11:32:05.678537Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Define CycleGan model","metadata":{}},{"cell_type":"code","source":"class CycleGan(tf.keras.Model):\n    def __init__(self, G_A2B, G_B2A, D_B, D_A):\n        super(CycleGan, self).__init__()\n        self.G_A2B = G_A2B\n        self.G_B2A = G_B2A\n        self.D_B = D_B\n        self.D_A = D_A\n      \n    def compile(\n            self, \n            G_optimizer, \n            D_optimizer,\n            disc_loss_fn, \n            gen_loss_fn, \n            cycle_loss_fn, \n            identity_loss_fn\n        ):\n        super(CycleGan, self).compile()\n        self.G_optimizer = G_optimizer\n        self.D_optimizer = D_optimizer\n        self.d_loss_fn = disc_loss_fn\n        self.g_loss_fn = gen_loss_fn\n        self.cycle_loss_fn = cycle_loss_fn\n        self.identity_loss_fn = identity_loss_fn\n        \n    def call(self, inputs):\n        return self.G_B2A(inputs)\n    \n    def train_step(self, input_batch):\n        A, B = input_batch\n        with tf.GradientTape(persistent=True) as tape:\n            A2B = self.G_A2B(A, training=True)\n            B2A = self.G_B2A(B, training=True)\n            \n            A2B2A = self.G_B2A(A2B, training=True)\n            B2A2B = self.G_A2B(B2A, training=True)\n            \n            A2A = self.G_B2A(A, training=True)\n            B2B = self.G_A2B(B, training=True)\n\n            A2B_d_logits = self.D_B(A2B, training=True)\n            B2A_d_logits = self.D_A(B2A, training=True)\n\n            A2B_g_loss = self.g_loss_fn(A2B_d_logits)\n            B2A_g_loss = self.g_loss_fn(B2A_d_logits)\n            \n            A2B2A_cycle_loss = self.cycle_loss_fn(A, A2B2A)\n            B2A2B_cycle_loss = self.cycle_loss_fn(B, B2A2B)\n            \n            A2A_id_loss = self.identity_loss_fn(A, A2A)\n            B2B_id_loss = self.identity_loss_fn(B, B2B)\n            \n            G_loss = (A2B_g_loss + B2A_g_loss) + (A2B2A_cycle_loss + B2A2B_cycle_loss) * 10.0 + (A2A_id_loss + B2B_id_loss) * 0.0\n           \n#             A2B = A2B_pool(A2B)  # or A2B = A2B_pool(A2B.numpy()), but it is much slower\n#             B2A = B2A_pool(B2A)  # because of the communication between CPU and GPU\n            \n            A_d_logits = self.D_A(A, training=True)\n            B2A_d_logits = self.D_A(B2A, training=True)\n            \n            B_d_logits = self.D_B(B, training=True)\n            A2B_d_logits = self.D_B(A2B, training=True)\n\n            A_d_loss, B2A_d_loss = self.d_loss_fn(A_d_logits, B2A_d_logits)\n            B_d_loss, A2B_d_loss = self.d_loss_fn(B_d_logits, A2B_d_logits)\n            \n            D_A_gp = gradient_penalty(functools.partial(self.D_A, training=True), A, B2A, mode='none')\n            D_B_gp = gradient_penalty(functools.partial(self.D_B, training=True), B, A2B, mode='none')\n\n            D_loss = (A_d_loss + B2A_d_loss) + (B_d_loss + A2B_d_loss) + (D_A_gp + D_B_gp) * 10.0\n\n        G_grad = tape.gradient(G_loss, self.G_A2B.trainable_variables + self.G_B2A.trainable_variables)\n        self.G_optimizer.apply_gradients(zip(G_grad,self. G_A2B.trainable_variables + self.G_B2A.trainable_variables))\n        \n        D_grad = tape.gradient(D_loss, self.D_A.trainable_variables + self.D_B.trainable_variables)\n        self.D_optimizer.apply_gradients(zip(D_grad, self.D_A.trainable_variables + self.D_B.trainable_variables))\n\n        return {\n            \"monet_gen_loss\": B2A_g_loss,\n            \"photo_gen_loss\": A2B_g_loss,\n            \"G_gen_loss\": G_loss,\n            \"monet_disc_loss\": A_d_loss,\n            \"photo_disc_loss\": B_d_loss,\n            \"D_disc_loss\": D_loss\n        }  ","metadata":{"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2021-10-05T11:32:05.680603Z","iopub.execute_input":"2021-10-05T11:32:05.681086Z","iopub.status.idle":"2021-10-05T11:32:05.698892Z","shell.execute_reply.started":"2021-10-05T11:32:05.681046Z","shell.execute_reply":"2021-10-05T11:32:05.698031Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Create and compile models, define losses and optimizers","metadata":{}},{"cell_type":"code","source":"# with strategy.scope(): \nG_A2B=ResnetGenerator(input_shape=SHAPE, norm='batch_norm') # layer_norm batch_norm instance_norm\nG_B2A=ResnetGenerator(input_shape=SHAPE, norm='batch_norm') # layer_norm batch_norm instance_norm\nD_B=ConvDiscriminator(input_shape=SHAPE, norm='batch_norm') # layer_norm batch_norm instance_norm\nD_A=ConvDiscriminator(input_shape=SHAPE, norm='batch_norm') # layer_norm batch_norm instance_norm\n\nmodel = CycleGan(G_A2B=G_A2B,G_B2A=G_B2A,D_B=D_B,D_A=D_A)\n\nG_lr_scheduler = LinearDecay(START_LEARNING_RATE, EPOCHS * STEPS_PER_EPOCH, STEPS_PER_EPOCH)\nD_lr_scheduler = LinearDecay(START_LEARNING_RATE_DISC, EPOCHS * STEPS_PER_EPOCH, STEPS_PER_EPOCH)\n\nmse = tf.losses.MeanSquaredError()\n\ndef discriminator_loss(r_logit, f_logit):\n    r_loss = mse(tf.ones_like(r_logit), r_logit)\n    f_loss = mse(tf.zeros_like(f_logit), f_logit)\n    return r_loss, f_loss\n\ndef generator_loss(f_logit):\n    f_loss = mse(tf.ones_like(f_logit), f_logit)\n    return f_loss\n\n# bce = tf.losses.BinaryCrossentropy(from_logits=True)\n\n# def discriminator_loss(r_logit, f_logit):\n#     r_loss = bce(tf.ones_like(r_logit), r_logit)\n#     f_loss = bce(tf.zeros_like(f_logit), f_logit)\n#     return r_loss, f_loss\n\n# def generator_loss(f_logit):\n#     f_loss = bce(tf.ones_like(f_logit), f_logit)\n#     return f_loss\n\nmodel.compile(\n    G_optimizer=Adam(learning_rate=G_lr_scheduler, beta_1=0.5),\n    D_optimizer=Adam(learning_rate=D_lr_scheduler, beta_1=0.5),\n    disc_loss_fn=discriminator_loss,\n    gen_loss_fn=generator_loss,\n    cycle_loss_fn=tf.losses.MeanAbsoluteError(),\n    identity_loss_fn=tf.losses.MeanAbsoluteError()\n)","metadata":{"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2021-10-05T11:32:05.700273Z","iopub.execute_input":"2021-10-05T11:32:05.700683Z","iopub.status.idle":"2021-10-05T11:32:06.658335Z","shell.execute_reply.started":"2021-10-05T11:32:05.700650Z","shell.execute_reply":"2021-10-05T11:32:06.657486Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cp_callback = tf.keras.callbacks.ModelCheckpoint(\n    filepath=\"/kaggle/working/cp-last-gpu.h5\", \n    monitor='photo_gen_loss', \n    verbose=1, \n    mode='min', \n    save_weights_only=True,\n    save_freq=SAVE_FREQ\n)\n                                         \nclass DisplayCallback(tf.keras.callbacks.Callback):\n    def __init__(self):\n        super().__init__()\n        self._supports_tf_logs = True\n    def on_epoch_end(self, epoch, logs=None):\n        clear_output(wait=True)\n        monet, photo = next(iter(test_dataset))\n        display_epoch_predict(G_B2A.predict(photo), photo, monet, G_A2B.predict(monet))\n        print(f'm_gen lr: {self.model.G_optimizer._decayed_lr(tf.float32).numpy()}')\n        print(f'm_disc lr: {self.model.D_optimizer._decayed_lr(tf.float32).numpy()}')\n\nclass SaveModelsWeightsCallback(tf.keras.callbacks.Callback):\n    def __init__(self):\n        super().__init__()\n    def on_epoch_end(self, epoch, logs=None):\n        G_A2B.save_weights('./G_A2B-tpu.h5')\n        G_B2A.save_weights('./G_B2A-tpu.h5')\n\nplot_training = PlotTraining(sample_rate=8, zoom=12)                                         \ncallbacks_list = [DisplayCallback(), plot_training, cp_callback, SaveModelsWeightsCallback()]","metadata":{"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2021-10-05T11:32:06.659751Z","iopub.execute_input":"2021-10-05T11:32:06.660098Z","iopub.status.idle":"2021-10-05T11:32:06.669594Z","shell.execute_reply.started":"2021-10-05T11:32:06.660063Z","shell.execute_reply":"2021-10-05T11:32:06.668403Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model.built = True\nmodel.predict(next(iter(dataset))[0])\nmodel.load_weights('../input/checkpoints/cp-last-gpu.h5')","metadata":{"execution":{"iopub.status.busy":"2021-10-05T11:32:06.671240Z","iopub.execute_input":"2021-10-05T11:32:06.671678Z","iopub.status.idle":"2021-10-05T11:32:12.999827Z","shell.execute_reply.started":"2021-10-05T11:32:06.671642Z","shell.execute_reply":"2021-10-05T11:32:12.998724Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model.fit(dataset, epochs=EPOCHS, steps_per_epoch=STEPS_PER_EPOCH, callbacks=callbacks_list)","metadata":{"execution":{"iopub.status.busy":"2021-10-05T11:32:13.001558Z","iopub.execute_input":"2021-10-05T11:32:13.001911Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model.save_weights('./cp-last-gpu.h5')\nG_B2A.save_weights('./G_B2A.h5')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# FIT NATIVE EXAMPLE","metadata":{}},{"cell_type":"code","source":"# with strategy.scope():\n#     A2B_pool = ItemPool(50)\n#     B2A_pool = ItemPool(50)\n\n# class CycleGan(tf.keras.Model):\n#     def __init__(self, G_A2B, G_B2A, D_B, D_A):\n#         super(CycleGan, self).__init__()\n#         self.G_A2B = G_A2B\n#         self.G_B2A = G_B2A\n#         self.D_B = D_B\n#         self.D_A = D_A\n      \n#     def compile(\n#             self, \n#             G_optimizer, \n#             D_optimizer,\n#             disc_loss_fn, \n#             gen_loss_fn, \n#             cycle_loss_fn, \n#             identity_loss_fn\n#         ):\n#         super(CycleGan, self).compile()\n#         self.G_optimizer = G_optimizer\n#         self.D_optimizer = D_optimizer\n#         self.d_loss_fn = disc_loss_fn\n#         self.g_loss_fn = gen_loss_fn\n#         self.cycle_loss_fn = cycle_loss_fn\n#         self.identity_loss_fn = identity_loss_fn\n        \n#     def call(self, inputs):\n#         return self.G_B2A(inputs)\n    \n#     def train_step(self, input_batch):\n#         A, B = input_batch\n#         with tf.GradientTape(persistent=True) as tape:\n#             A2B = self.G_A2B(A, training=True)\n#             B2A = self.G_B2A(B, training=True)\n            \n#             A2B2A = self.G_B2A(A2B, training=True)\n#             B2A2B = self.G_A2B(B2A, training=True)\n            \n#             A2A = self.G_B2A(A, training=True)\n#             B2B = self.G_A2B(B, training=True)\n\n#             A2B_d_logits = self.D_B(A2B, training=True)\n#             B2A_d_logits = self.D_A(B2A, training=True)\n\n#             A2B_g_loss = self.g_loss_fn(A2B_d_logits)\n#             B2A_g_loss = self.g_loss_fn(B2A_d_logits)\n            \n#             A2B2A_cycle_loss = self.cycle_loss_fn(A, A2B2A)\n#             B2A2B_cycle_loss = self.cycle_loss_fn(B, B2A2B)\n            \n#             A2A_id_loss = self.identity_loss_fn(A, A2A)\n#             B2B_id_loss = self.identity_loss_fn(B, B2B)\n            \n#             G_loss = (A2B_g_loss + B2A_g_loss) + (A2B2A_cycle_loss + B2A2B_cycle_loss) * 10.0 + (A2A_id_loss + B2B_id_loss) * 0.0\n           \n# #             A2B = A2B_pool(A2B)  # or A2B = A2B_pool(A2B.numpy()), but it is much slower\n# #             B2A = B2A_pool(B2A)  # because of the communication between CPU and GPU\n            \n#             A_d_logits = self.D_A(A, training=True)\n#             B2A_d_logits = self.D_A(B2A, training=True)\n            \n#             B_d_logits = self.D_B(B, training=True)\n#             A2B_d_logits = self.D_B(A2B, training=True)\n\n#             A_d_loss, B2A_d_loss = self.d_loss_fn(A_d_logits, B2A_d_logits)\n#             B_d_loss, A2B_d_loss = self.d_loss_fn(B_d_logits, A2B_d_logits)\n            \n#             D_A_gp = gradient_penalty(functools.partial(self.D_A, training=True), A, B2A, mode='none')\n#             D_B_gp = gradient_penalty(functools.partial(self.D_B, training=True), B, A2B, mode='none')\n\n#             D_loss = (A_d_loss + B2A_d_loss) + (B_d_loss + A2B_d_loss) + (D_A_gp + D_B_gp) * 10.0\n\n#         G_grad = tape.gradient(G_loss, self.G_A2B.trainable_variables + self.G_B2A.trainable_variables)\n#         self.G_optimizer.apply_gradients(zip(G_grad,self. G_A2B.trainable_variables + self.G_B2A.trainable_variables))\n        \n#         D_grad = tape.gradient(D_loss, self.D_A.trainable_variables + self.D_B.trainable_variables)\n#         self.D_optimizer.apply_gradients(zip(D_grad, self.D_A.trainable_variables + self.D_B.trainable_variables))\n\n#         return {\n#             \"monet_gen_loss\": B2A_g_loss,\n#             \"photo_gen_loss\": A2B_g_loss,\n#             \"G_gen_loss\": G_loss,\n#             \"monet_disc_loss\": A_d_loss,\n#             \"photo_disc_loss\": B_d_loss,\n#             \"D_disc_loss\": D_loss\n#         }\n                                     \n# # with strategy.scope():    \n# model = CycleGan(\n#     G_A2B=ResnetGenerator(input_shape=SHAPE),\n#     G_B2A=ResnetGenerator(input_shape=SHAPE),\n#     D_B=ConvDiscriminator(input_shape=SHAPE),\n#     D_A=ConvDiscriminator(input_shape=SHAPE)\n# )\n\n# G_lr_scheduler = LinearDecay(0.0002, EPOCHS * STEPS_PER_EPOCH, STEPS_PER_EPOCH)\n# D_lr_scheduler = LinearDecay(0.0002, EPOCHS * STEPS_PER_EPOCH, STEPS_PER_EPOCH)\n\n# mse = tf.losses.MeanSquaredError()\n\n# def discriminator_loss(r_logit, f_logit):\n#     r_loss = mse(tf.ones_like(r_logit), r_logit)\n#     f_loss = mse(tf.zeros_like(f_logit), f_logit)\n#     return r_loss, f_loss\n\n# def generator_loss(f_logit):\n#     f_loss = mse(tf.ones_like(f_logit), f_logit)\n#     return f_loss\n\n# model.compile(\n#     G_optimizer=Adam(learning_rate=G_lr_scheduler, beta_1=0.5),\n#     D_optimizer=Adam(learning_rate=D_lr_scheduler, beta_1=0.5),\n#     disc_loss_fn=discriminator_loss,\n#     gen_loss_fn=generator_loss,\n#     cycle_loss_fn=tf.losses.MeanAbsoluteError(),\n#     identity_loss_fn=tf.losses.MeanAbsoluteError()\n# )\n\n# cp_callback = tf.keras.callbacks.ModelCheckpoint(\n#     filepath=\"/kaggle/working/cp-last.h5\", \n#     monitor='monet_gen_loss', \n#     verbose=1, \n#     mode='min', \n#     save_weights_only=True,\n#     save_freq=SAVE_FREQ\n# )\n                                         \n# class DisplayCallback(tf.keras.callbacks.Callback):\n#     def __init__(self):\n#         super().__init__()\n#         self._supports_tf_logs = True\n#     def on_epoch_end(self, epoch, logs=None):\n#         clear_output(wait=True)\n#         simple_image = next(iter(test_photo))\n#         plot_images(self.model.predict(simple_image), simple_image)\n#         print(f'm_gen lr: {self.model.G_optimizer._decayed_lr(tf.float32).numpy()}')\n#         print(f'm_disc lr: {self.model.D_optimizer._decayed_lr(tf.float32).numpy()}')\n\n# plot_training = PlotTraining(sample_rate=8, zoom=12)                                         \n# callbacks_list = [DisplayCallback(), plot_training, cp_callback] # cp_callback, plot_training \n                                         \n# model.fit(dataset, epochs=EPOCHS, steps_per_epoch=STEPS_PER_EPOCH, callbacks=callbacks_list)","metadata":{"jupyter":{"source_hidden":true},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# CUSTOM NATIVE EXAMPLE","metadata":{}},{"cell_type":"code","source":"# args = {\n#     'dataset': 'horse2zebra',\n#     'datasets_dir': '../input/gan-getting-started',\n#     'load_size': 286,\n#     'crop_size': 256,\n#     'batch_size': 1,\n#     'epochs':EPOCHS,\n#     'epoch_decay':EPOCHS,\n#     'lr':0.0002,\n#     'beta_1':0.5,\n#     'adversarial_loss_mode':'lsgan',\n#     'gradient_penalty_mode':'none',\n#     'gradient_penalty_weight':10.0,\n#     'cycle_loss_weight':10.0,\n#     'identity_loss_weight':0.0,\n#     'pool_size':50,\n# }\n# # output_dir\n# output_dir = 'output_dir/'\n# !mkdir output_dir\n\n# A_img_paths = tf.io.gfile.glob(GCS_DS_PATH+'monet_jpg/*.jpg')\n# B_img_paths = tf.io.gfile.glob(GCS_DS_PATH+'photo_jpg/*.jpg')\n# A_B_dataset, len_dataset = make_zip_dataset(A_img_paths, B_img_paths, BATCH_SIZE, 286, 256, training=True, repeat=False)\n# A_B_dataset = dataset\n# A2B_pool = ItemPool(50)\n# B2A_pool = ItemPool(50)\n\n# G_A2B = ResnetGenerator(input_shape=SHAPE)\n# G_B2A = ResnetGenerator(input_shape=SHAPE)\n\n# D_A = ConvDiscriminator(input_shape=SHAPE)\n# D_B = ConvDiscriminator(input_shape=SHAPE)\n\n# d_loss_fn, g_loss_fn = get_adversarial_losses_fn('lsgan')\n# cycle_loss_fn = tf.losses.MeanAbsoluteError()\n# identity_loss_fn = tf.losses.MeanAbsoluteError()\n# # lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(\n# #     0.0004, decay_steps=STEPS_PER_EPOCH, decay_rate=0.85, staircase=False, name=None\n# # )\n# G_lr_scheduler = LinearDecay(0.0002, EPOCHS * STEPS_PER_EPOCH, STEPS_PER_EPOCH)\n# D_lr_scheduler = LinearDecay(0.0002, EPOCHS * STEPS_PER_EPOCH, STEPS_PER_EPOCH)\n# G_optimizer = Adam(learning_rate=G_lr_scheduler, beta_1=0.5)\n# D_optimizer = Adam(learning_rate=D_lr_scheduler, beta_1=0.5)\n\n# @tf.function\n# def train_G(A, B):\n#     with tf.GradientTape() as t:\n#         A2B = G_A2B(A, training=True)\n#         B2A = G_B2A(B, training=True)\n#         A2B2A = G_B2A(A2B, training=True)\n#         B2A2B = G_A2B(B2A, training=True)\n#         A2A = G_B2A(A, training=True)\n#         B2B = G_A2B(B, training=True)\n\n#         A2B_d_logits = D_B(A2B, training=True)\n#         B2A_d_logits = D_A(B2A, training=True)\n\n#         A2B_g_loss = g_loss_fn(A2B_d_logits)\n#         B2A_g_loss = g_loss_fn(B2A_d_logits)\n#         A2B2A_cycle_loss = cycle_loss_fn(A, A2B2A)\n#         B2A2B_cycle_loss = cycle_loss_fn(B, B2A2B)\n#         A2A_id_loss = identity_loss_fn(A, A2A)\n#         B2B_id_loss = identity_loss_fn(B, B2B)\n\n#         G_loss = (A2B_g_loss + B2A_g_loss) + (A2B2A_cycle_loss + B2A2B_cycle_loss) * 10.0 + (A2A_id_loss + B2B_id_loss) * 0.0\n\n#     G_grad = t.gradient(G_loss, G_A2B.trainable_variables + G_B2A.trainable_variables)\n#     G_optimizer.apply_gradients(zip(G_grad, G_A2B.trainable_variables + G_B2A.trainable_variables))\n\n#     return A2B, B2A, {'A2B_g_loss': A2B_g_loss,\n#                       'B2A_g_loss': B2A_g_loss,\n#                       'A2B2A_cycle_loss': A2B2A_cycle_loss,\n#                       'B2A2B_cycle_loss': B2A2B_cycle_loss,\n#                       'A2A_id_loss': A2A_id_loss,\n#                       'B2B_id_loss': B2B_id_loss}\n\n\n# @tf.function\n# def train_D(A, B, A2B, B2A):\n#     with tf.GradientTape() as t:\n#         A_d_logits = D_A(A, training=True)\n#         B2A_d_logits = D_A(B2A, training=True)\n#         B_d_logits = D_B(B, training=True)\n#         A2B_d_logits = D_B(A2B, training=True)\n\n#         A_d_loss, B2A_d_loss = d_loss_fn(A_d_logits, B2A_d_logits)\n#         B_d_loss, A2B_d_loss = d_loss_fn(B_d_logits, A2B_d_logits)\n#         D_A_gp = gradient_penalty(functools.partial(D_A, training=True), A, B2A, mode='none')\n#         D_B_gp = gradient_penalty(functools.partial(D_B, training=True), B, A2B, mode='none')\n\n#         D_loss = (A_d_loss + B2A_d_loss) + (B_d_loss + A2B_d_loss) + (D_A_gp + D_B_gp) * 10.0\n\n#     D_grad = t.gradient(D_loss, D_A.trainable_variables + D_B.trainable_variables)\n#     D_optimizer.apply_gradients(zip(D_grad, D_A.trainable_variables + D_B.trainable_variables))\n\n#     return {'A_d_loss': A_d_loss + B2A_d_loss,\n#             'B_d_loss': B_d_loss + A2B_d_loss,\n#             'D_A_gp': D_A_gp,\n#             'D_B_gp': D_B_gp}\n\n\n# def train_step(A, B):\n#     A2B, B2A, G_loss_dict = train_G(A, B)\n\n#     # cannot autograph `A2B_pool`\n#     A2B = A2B_pool(A2B)  # or A2B = A2B_pool(A2B.numpy()), but it is much slower\n#     B2A = B2A_pool(B2A)  # because of the communication between CPU and GPU\n\n#     D_loss_dict = train_D(A, B, A2B, B2A)\n\n#     return G_loss_dict, D_loss_dict\n\n\n# @tf.function\n# def sample(A, B):\n#     A2B = G_A2B(A, training=False)\n#     B2A = G_B2A(B, training=False)\n#     A2B2A = G_B2A(A2B, training=False)\n#     B2A2B = G_A2B(B2A, training=False)\n#     return A2B, B2A, A2B2A, B2A2B\n\n# ep_cnt = tf.Variable(initial_value=0, trainable=False, dtype=tf.int64)\n\n# # checkpoint\n# # summary\n# train_summary_writer = tf.summary.create_file_writer(output_dir+ 'summaries'+'/train')\n\n# # sample\n# sample_dir = output_dir + 'samples_training/'\n# !mkdir sample_dir\n\n# data_iter = iter(A_B_dataset)\n# # main loop\n# with train_summary_writer.as_default():\n#     for ep in range(EPOCHS):\n#         start_epoch_time = time.time()\n#         if ep < ep_cnt:\n#             continue\n\n#         # update epoch counter\n#         ep_cnt.assign_add(1)\n\n#         # train for an epoch\n#         for step in range(STEPS_PER_EPOCH):\n#             A, B = next(data_iter)\n#             G_loss_dict, D_loss_dict = train_step(A, B)\n#             del A, B\n#             gc.collect()\n#             if (step+1) % STEPS_PER_EPOCH == 0:\n#                 clear_output(wait=True)\n#                 print ('Time taken for Epoch {}/{} is {} \\n'.format(ep+1, EPOCHS, str(datetime.timedelta(seconds=int(time.time()-start_epoch_time)))))\n#                 print(f'm_gen lr: {G_optimizer._decayed_lr(tf.float32).numpy()}')\n#                 print(f'm_disc lr: {D_optimizer._decayed_lr(tf.float32).numpy()} \\n')\n#                 print('A2B_g_loss : ' , G_loss_dict['A2B_g_loss'].numpy()) \n#                 print('B2A_g_loss : ' , G_loss_dict['B2A_g_loss'].numpy()) \n#                 print('A2B2A_cycle_loss : ' , G_loss_dict['A2B2A_cycle_loss'].numpy()) \n#                 print('B2A2B_cycle_loss : ' ,  G_loss_dict['B2A2B_cycle_loss'].numpy()) \n#                 print('A2A_id_loss : ' , G_loss_dict['A2A_id_loss'].numpy()) \n#                 print('B2B_id_loss : ' , G_loss_dict['B2B_id_loss'].numpy())\n#                 print('A_d_loss : ' , D_loss_dict['A_d_loss'].numpy())\n#                 print('B_d_loss : ' , D_loss_dict['B_d_loss'].numpy())\n#                 print('D_A_gp : ' , D_loss_dict['D_A_gp'].numpy())\n#                 print('D_B_gp : ' , D_loss_dict['D_B_gp'].numpy())\n#                 A, B = next(iter(dataset))\n#                 A2B, B2A, A2B2A, B2A2B = sample(A, B)\n#                 plot_images(B2A, B)\n#                 del A\n#                 del B\n#                 del A2B, B2A, A2B2A, B2A2B\n#                 gc.collect()","metadata":{"jupyter":{"source_hidden":true},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for inp in test_photo.take(10):\n    plot_images(model.predict(inp), inp)","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}