{"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":"code","source":"from pathlib import Path\n\nimport tensorflow as tf\nfrom tensorflow import keras\nfrom tensorflow.keras import layers\nimport tensorflow_addons as tfa\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:04.692220Z","iopub.execute_input":"2022-11-18T18:32:04.693005Z","iopub.status.idle":"2022-11-18T18:32:09.727550Z","shell.execute_reply.started":"2022-11-18T18:32:04.692906Z","shell.execute_reply":"2022-11-18T18:32:09.726772Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!nvidia-smi","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:09.730182Z","iopub.execute_input":"2022-11-18T18:32:09.730715Z","iopub.status.idle":"2022-11-18T18:32:10.736684Z","shell.execute_reply.started":"2022-11-18T18:32:09.730677Z","shell.execute_reply":"2022-11-18T18:32:10.735807Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Enable TPU if needed.","metadata":{}},{"cell_type":"code","source":"# Not using TPU\n# try:\n#     raise ValueError('disable TPU')\n#     tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n#     print('Device:', tpu.master())\n#     tf.config.experimental_connect_to_cluster(tpu)\n#     tf.tpu.experimental.initialize_tpu_system(tpu)\n#     strategy = tf.distribute.experimental.TPUStrategy(tpu)\n# except:\n#     print('...fail...')\n#     strategy = tf.distribute.get_strategy()\n# print('Number of replicas:', strategy.num_replicas_in_sync)\n\nstrategy = tf.distribute.get_strategy()\nAUTOTUNE = tf.data.experimental.AUTOTUNE","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:10.738554Z","iopub.execute_input":"2022-11-18T18:32:10.738893Z","iopub.status.idle":"2022-11-18T18:32:10.748814Z","shell.execute_reply.started":"2022-11-18T18:32:10.738838Z","shell.execute_reply":"2022-11-18T18:32:10.748061Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"tf.__version__","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:10.750466Z","iopub.execute_input":"2022-11-18T18:32:10.750994Z","iopub.status.idle":"2022-11-18T18:32:10.763766Z","shell.execute_reply.started":"2022-11-18T18:32:10.750950Z","shell.execute_reply":"2022-11-18T18:32:10.762818Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Sketch Dataset","metadata":{}},{"cell_type":"code","source":"img_height, img_width = 256, 256\nchannels = 3","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:10.766494Z","iopub.execute_input":"2022-11-18T18:32:10.766821Z","iopub.status.idle":"2022-11-18T18:32:10.772734Z","shell.execute_reply.started":"2022-11-18T18:32:10.766788Z","shell.execute_reply":"2022-11-18T18:32:10.771837Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def decode_img(image):\n    # convert to 3-channel\n    image = tf.image.decode_jpeg(image, channels=3)\n    image = tf.image.resize(image, [img_height, img_width])\n    # scale to <-1,1>\n    image = (tf.cast(image, tf.float32) / 127.5) - 1\n    return image\n\ndef process_path(file_path):\n    # does not work when using TPU\n    img = tf.io.read_file(file_path)\n    img = decode_img(img)\n    return img","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:10.774498Z","iopub.execute_input":"2022-11-18T18:32:10.775101Z","iopub.status.idle":"2022-11-18T18:32:10.784092Z","shell.execute_reply.started":"2022-11-18T18:32:10.775065Z","shell.execute_reply":"2022-11-18T18:32:10.783264Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"list_ds = tf.data.Dataset.list_files('../input/imagenetsketch/sketch/*/*.JPEG', shuffle=True)\nstyled_ds = list_ds.map(process_path, num_parallel_calls=tf.data.experimental.AUTOTUNE).batch(1)","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:10.785466Z","iopub.execute_input":"2022-11-18T18:32:10.785806Z","iopub.status.idle":"2022-11-18T18:32:17.896351Z","shell.execute_reply.started":"2022-11-18T18:32:10.785775Z","shell.execute_reply":"2022-11-18T18:32:17.895582Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def plot(image, ax=None):\n    if ax is None:\n        ax = plt.gca()\n    ax.imshow(image * 0.5 + 0.5)","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:17.897636Z","iopub.execute_input":"2022-11-18T18:32:17.897907Z","iopub.status.idle":"2022-11-18T18:32:17.904552Z","shell.execute_reply.started":"2022-11-18T18:32:17.897855Z","shell.execute_reply":"2022-11-18T18:32:17.903852Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fix, axs = plt.subplots(1, 5, figsize=(17,17))\nfor batch, ax in zip(styled_ds.take(5), axs):\n    plot(batch[0], ax)\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:17.905572Z","iopub.execute_input":"2022-11-18T18:32:17.905835Z","iopub.status.idle":"2022-11-18T18:32:18.677928Z","shell.execute_reply.started":"2022-11-18T18:32:17.905800Z","shell.execute_reply":"2022-11-18T18:32:18.677228Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"list_ds = tf.data.Dataset.list_files('../input/gan-getting-started/photo_jpg/*.jpg', shuffle=True)\nphoto_ds = list_ds.map(process_path, num_parallel_calls=tf.data.experimental.AUTOTUNE).batch(1)","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:18.678972Z","iopub.execute_input":"2022-11-18T18:32:18.679235Z","iopub.status.idle":"2022-11-18T18:32:24.402323Z","shell.execute_reply.started":"2022-11-18T18:32:18.679203Z","shell.execute_reply":"2022-11-18T18:32:24.401546Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fix, axs = plt.subplots(1, 5, figsize=(17,17))\nfor batch, ax in zip(photo_ds.take(5), axs):\n    plot(batch[0], ax)\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:24.403571Z","iopub.execute_input":"2022-11-18T18:32:24.403823Z","iopub.status.idle":"2022-11-18T18:32:25.050387Z","shell.execute_reply.started":"2022-11-18T18:32:24.403792Z","shell.execute_reply":"2022-11-18T18:32:25.049773Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"OUTPUT_CHANNELS = 3\n\ndef downsample(filters, size, apply_instancenorm=True):\n    initializer = tf.random_normal_initializer(0., 0.02)\n    gamma_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n\n    result = keras.Sequential()\n    result.add(layers.Conv2D(filters, size, strides=2, padding='same',\n                             kernel_initializer=initializer, use_bias=False))\n\n    if apply_instancenorm:\n        result.add(tfa.layers.InstanceNormalization(gamma_initializer=gamma_init))\n\n    result.add(layers.LeakyReLU())\n\n    return result","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:25.052680Z","iopub.execute_input":"2022-11-18T18:32:25.053203Z","iopub.status.idle":"2022-11-18T18:32:25.060376Z","shell.execute_reply.started":"2022-11-18T18:32:25.053168Z","shell.execute_reply":"2022-11-18T18:32:25.059446Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def upsample(filters, size, apply_dropout=False):\n    initializer = tf.random_normal_initializer(0., 0.02)\n    gamma_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n\n    result = keras.Sequential()\n    result.add(layers.Conv2DTranspose(filters, size, strides=2,\n                                      padding='same',\n                                      kernel_initializer=initializer,\n                                      use_bias=False))\n\n    result.add(tfa.layers.InstanceNormalization(gamma_initializer=gamma_init))\n\n    if apply_dropout:\n        result.add(layers.Dropout(0.5))\n\n    result.add(layers.ReLU())\n\n    return result","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:25.061501Z","iopub.execute_input":"2022-11-18T18:32:25.062305Z","iopub.status.idle":"2022-11-18T18:32:25.072790Z","shell.execute_reply.started":"2022-11-18T18:32:25.062271Z","shell.execute_reply":"2022-11-18T18:32:25.072051Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def Generator():\n    inputs = layers.Input(shape=[img_height,img_width,3])\n\n    # bs = batch size, assuming (h, w) = (256, 256)\n    down_stack = [\n        downsample(64, 4, apply_instancenorm=False), \n        downsample(128, 4), \n        downsample(256, 4),\n        downsample(512, 4),\n        downsample(512, 4),\n        downsample(512, 4),\n        downsample(512, 4),\n        downsample(512, 4)\n    ]\n\n    up_stack = [\n        upsample(512, 4, apply_dropout=True), \n        upsample(512, 4, apply_dropout=True),\n        upsample(512, 4, apply_dropout=True),\n        upsample(512, 4), \n        upsample(256, 4),\n        upsample(128, 4), \n        upsample(64, 4), \n    ]\n\n    initializer = tf.random_normal_initializer(0., 0.02)\n    last = layers.Conv2DTranspose(OUTPUT_CHANNELS, 4,\n                                  strides=2,\n                                  padding='same',\n                                  kernel_initializer=initializer,\n                                  activation='tanh') # (bs, 256, 256, 3)\n\n    x = inputs\n\n    skips = []\n    for down in down_stack:\n        x = down(x)\n        skips.append(x)\n\n    skips = reversed(skips[:-1])\n\n  \n    for up, skip in zip(up_stack, skips):\n        x = up(x)\n        x = layers.Concatenate()([x, skip])\n\n    x = last(x)\n\n    return keras.Model(inputs=inputs, outputs=x)","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:25.077428Z","iopub.execute_input":"2022-11-18T18:32:25.077693Z","iopub.status.idle":"2022-11-18T18:32:25.260802Z","shell.execute_reply.started":"2022-11-18T18:32:25.077665Z","shell.execute_reply":"2022-11-18T18:32:25.259957Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def Discriminator():\n    initializer = tf.random_normal_initializer(0., 0.02)\n    gamma_init = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)\n\n    inp = layers.Input(shape=[256, 256, 3], name='input_image')\n\n    x = inp\n\n    down1 = downsample(64, 4, False)(x) \n    down2 = downsample(128, 4)(down1) \n    down3 = downsample(256, 4)(down2) \n\n    zero_pad1 = layers.ZeroPadding2D()(down3) \n    conv = layers.Conv2D(512, 4, strides=1,\n                         kernel_initializer=initializer,\n                         use_bias=False)(zero_pad1)\n\n    norm1 = tfa.layers.InstanceNormalization(gamma_initializer=gamma_init)(conv)\n\n    leaky_relu = layers.LeakyReLU()(norm1)\n\n    zero_pad2 = layers.ZeroPadding2D()(leaky_relu)\n\n    last = layers.Conv2D(1, 4, strides=1,\n                         kernel_initializer=initializer)(zero_pad2)\n\n    return tf.keras.Model(inputs=inp, outputs=last)","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:25.262072Z","iopub.execute_input":"2022-11-18T18:32:25.262473Z","iopub.status.idle":"2022-11-18T18:32:25.277255Z","shell.execute_reply.started":"2022-11-18T18:32:25.262438Z","shell.execute_reply":"2022-11-18T18:32:25.276499Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with strategy.scope():\n    styled_generator = Generator() #  photos to sketch\n    photo_generator = Generator() #  sketch to photos\n\n    styled_discriminator = Discriminator() # differentiates real and generated sketch\n    photo_discriminator = Discriminator() # differentiates real and generated photos","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:25.279815Z","iopub.execute_input":"2022-11-18T18:32:25.280135Z","iopub.status.idle":"2022-11-18T18:32:27.934554Z","shell.execute_reply.started":"2022-11-18T18:32:25.280102Z","shell.execute_reply":"2022-11-18T18:32:27.933814Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"example = next(iter(photo_ds))","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:27.938652Z","iopub.execute_input":"2022-11-18T18:32:27.939488Z","iopub.status.idle":"2022-11-18T18:32:27.986216Z","shell.execute_reply.started":"2022-11-18T18:32:27.939452Z","shell.execute_reply":"2022-11-18T18:32:27.985522Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"example.shape","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:27.987686Z","iopub.execute_input":"2022-11-18T18:32:27.988184Z","iopub.status.idle":"2022-11-18T18:32:27.994425Z","shell.execute_reply.started":"2022-11-18T18:32:27.988150Z","shell.execute_reply":"2022-11-18T18:32:27.993596Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"styled = styled_generator(example)\n\nplt.subplot(1, 2, 1)\nplt.title(\"Original\")\nplot(example[0])\nplt.show()\n\nplt.subplot(1, 2, 2)\nplt.title(\"Generated\")\nplot(styled[0])\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:27.995836Z","iopub.execute_input":"2022-11-18T18:32:27.996398Z","iopub.status.idle":"2022-11-18T18:32:35.097632Z","shell.execute_reply.started":"2022-11-18T18:32:27.996364Z","shell.execute_reply":"2022-11-18T18:32:35.096818Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## CycleGAN","metadata":{}},{"cell_type":"code","source":"class CycleGan(keras.Model):\n    def __init__(\n        self,\n        styled_generator,\n        photo_generator,\n        styled_discriminator,\n        photo_discriminator,\n        lambda_cycle=10,\n    ):\n        super(CycleGan, self).__init__()\n        self.s_gen = styled_generator\n        self.p_gen = photo_generator\n        self.s_disc = styled_discriminator\n        self.p_disc = photo_discriminator\n        self.lambda_cycle = lambda_cycle\n        \n    def compile(\n        self,\n        s_gen_optimizer,\n        p_gen_optimizer,\n        s_disc_optimizer,\n        p_disc_optimizer,\n        gen_loss_fn,\n        disc_loss_fn,\n        cycle_loss_fn,\n        identity_loss_fn\n    ):\n        super(CycleGan, self).compile()\n        self.s_gen_optimizer = s_gen_optimizer\n        self.p_gen_optimizer = p_gen_optimizer\n        self.s_disc_optimizer = s_disc_optimizer\n        self.p_disc_optimizer = p_disc_optimizer\n        self.gen_loss_fn = gen_loss_fn\n        self.disc_loss_fn = disc_loss_fn\n        self.cycle_loss_fn = cycle_loss_fn\n        self.identity_loss_fn = identity_loss_fn\n        \n    def train_step(self, batch_data):\n        real_styled, real_photo = batch_data\n        \n        with tf.GradientTape(persistent=True) as tape:\n            # photo to styled back to photo\n            fake_styled = self.s_gen(real_photo, training=True)\n            cycled_photo = self.p_gen(fake_styled, training=True)\n\n            # styled to photo back to styled\n            fake_photo = self.p_gen(real_styled, training=True)\n            cycled_styled = self.s_gen(fake_photo, training=True)\n\n            # generating itself\n            same_styled = self.s_gen(real_styled, training=True)\n            same_photo = self.p_gen(real_photo, training=True)\n\n            # discriminator used to check, inputing real images\n            disc_real_styled = self.s_disc(real_styled, training=True)\n            disc_real_photo = self.p_disc(real_photo, training=True)\n\n            # discriminator used to check, inputing fake images\n            disc_fake_styled = self.s_disc(fake_styled, training=True)\n            disc_fake_photo = self.p_disc(fake_photo, training=True)\n\n            # evaluates generator loss\n            styled_gen_loss = self.gen_loss_fn(disc_fake_styled)\n            photo_gen_loss = self.gen_loss_fn(disc_fake_photo)\n\n            # evaluates total cycle consistency loss\n            total_cycle_loss = self.cycle_loss_fn(real_styled, cycled_styled, self.lambda_cycle) + self.cycle_loss_fn(real_photo, cycled_photo, self.lambda_cycle)\n\n            # evaluates total generator loss\n            total_styled_gen_loss = styled_gen_loss + total_cycle_loss + self.identity_loss_fn(real_styled, same_styled, self.lambda_cycle)\n            total_photo_gen_loss = photo_gen_loss + total_cycle_loss + self.identity_loss_fn(real_photo, same_photo, self.lambda_cycle)\n\n            # evaluates discriminator loss\n            styled_disc_loss = self.disc_loss_fn(disc_real_styled, disc_fake_styled)\n            photo_disc_loss = self.disc_loss_fn(disc_real_photo, disc_fake_photo)\n\n        # Calculate the gradients for generator and discriminator\n        styled_generator_gradients = tape.gradient(total_styled_gen_loss,\n                                                  self.s_gen.trainable_variables)\n        photo_generator_gradients = tape.gradient(total_photo_gen_loss,\n                                                  self.p_gen.trainable_variables)\n\n        styled_discriminator_gradients = tape.gradient(styled_disc_loss,\n                                                      self.s_disc.trainable_variables)\n        photo_discriminator_gradients = tape.gradient(photo_disc_loss,\n                                                      self.p_disc.trainable_variables)\n\n        # Apply the gradients to the optimizer\n        self.s_gen_optimizer.apply_gradients(zip(styled_generator_gradients,\n                                                 self.s_gen.trainable_variables))\n\n        self.p_gen_optimizer.apply_gradients(zip(photo_generator_gradients,\n                                                 self.p_gen.trainable_variables))\n\n        self.s_disc_optimizer.apply_gradients(zip(styled_discriminator_gradients,\n                                                  self.s_disc.trainable_variables))\n\n        self.p_disc_optimizer.apply_gradients(zip(photo_discriminator_gradients,\n                                                  self.p_disc.trainable_variables))\n        \n        return {\n            \"styled_gen_loss\": total_styled_gen_loss,\n            \"photo_gen_loss\": total_photo_gen_loss,\n            \"styled_disc_loss\": styled_disc_loss,\n            \"photo_disc_loss\": photo_disc_loss\n        }","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:35.099512Z","iopub.execute_input":"2022-11-18T18:32:35.099863Z","iopub.status.idle":"2022-11-18T18:32:35.116620Z","shell.execute_reply.started":"2022-11-18T18:32:35.099820Z","shell.execute_reply":"2022-11-18T18:32:35.115840Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Loss functions","metadata":{"execution":{"iopub.status.busy":"2021-11-25T16:00:37.357996Z","iopub.execute_input":"2021-11-25T16:00:37.35828Z","iopub.status.idle":"2021-11-25T16:00:37.362523Z","shell.execute_reply.started":"2021-11-25T16:00:37.358249Z","shell.execute_reply":"2021-11-25T16:00:37.361787Z"}}},{"cell_type":"code","source":"with strategy.scope():\n    def discriminator_loss(real, generated):\n        real_loss = tf.keras.losses.BinaryCrossentropy(from_logits=True, reduction=tf.keras.losses.Reduction.NONE)(tf.ones_like(real), real)\n\n        generated_loss = tf.keras.losses.BinaryCrossentropy(from_logits=True, reduction=tf.keras.losses.Reduction.NONE)(tf.zeros_like(generated), generated)\n\n        total_disc_loss = real_loss + generated_loss\n\n        return total_disc_loss * 0.5","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:35.117862Z","iopub.execute_input":"2022-11-18T18:32:35.118216Z","iopub.status.idle":"2022-11-18T18:32:35.132513Z","shell.execute_reply.started":"2022-11-18T18:32:35.118180Z","shell.execute_reply":"2022-11-18T18:32:35.131772Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with strategy.scope():\n    def generator_loss(generated):\n        return tf.keras.losses.BinaryCrossentropy(from_logits=True, reduction=tf.keras.losses.Reduction.NONE)(tf.ones_like(generated), generated)","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:35.133853Z","iopub.execute_input":"2022-11-18T18:32:35.134138Z","iopub.status.idle":"2022-11-18T18:32:35.144896Z","shell.execute_reply.started":"2022-11-18T18:32:35.134105Z","shell.execute_reply":"2022-11-18T18:32:35.144232Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with strategy.scope():\n    def calc_cycle_loss(real_image, cycled_image, LAMBDA):\n        loss1 = tf.reduce_mean(tf.abs(real_image - cycled_image))\n\n        return LAMBDA * loss1","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:35.146115Z","iopub.execute_input":"2022-11-18T18:32:35.146561Z","iopub.status.idle":"2022-11-18T18:32:35.156249Z","shell.execute_reply.started":"2022-11-18T18:32:35.146518Z","shell.execute_reply":"2022-11-18T18:32:35.155452Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with strategy.scope():\n    def identity_loss(real_image, same_image, LAMBDA):\n        loss = tf.reduce_mean(tf.abs(real_image - same_image))\n        return LAMBDA * 0.5 * loss","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:35.157579Z","iopub.execute_input":"2022-11-18T18:32:35.157842Z","iopub.status.idle":"2022-11-18T18:32:35.167054Z","shell.execute_reply.started":"2022-11-18T18:32:35.157807Z","shell.execute_reply":"2022-11-18T18:32:35.166297Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training","metadata":{}},{"cell_type":"code","source":"with strategy.scope():\n    styled_generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n    photo_generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n\n    styled_discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n    photo_discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:35.168367Z","iopub.execute_input":"2022-11-18T18:32:35.168621Z","iopub.status.idle":"2022-11-18T18:32:35.178518Z","shell.execute_reply.started":"2022-11-18T18:32:35.168590Z","shell.execute_reply":"2022-11-18T18:32:35.177798Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with strategy.scope():\n    cycle_gan_model = CycleGan(\n        styled_generator, photo_generator, styled_discriminator, photo_discriminator\n    )\n\n    cycle_gan_model.compile(\n        s_gen_optimizer = styled_generator_optimizer,\n        p_gen_optimizer = photo_generator_optimizer,\n        s_disc_optimizer = styled_discriminator_optimizer,\n        p_disc_optimizer = photo_discriminator_optimizer,\n        gen_loss_fn = generator_loss,\n        disc_loss_fn = discriminator_loss,\n        cycle_loss_fn = calc_cycle_loss,\n        identity_loss_fn = identity_loss\n    )","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:35.180080Z","iopub.execute_input":"2022-11-18T18:32:35.180268Z","iopub.status.idle":"2022-11-18T18:32:35.206864Z","shell.execute_reply.started":"2022-11-18T18:32:35.180247Z","shell.execute_reply":"2022-11-18T18:32:35.206188Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(photo_ds), len(styled_ds)","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:35.209580Z","iopub.execute_input":"2022-11-18T18:32:35.209777Z","iopub.status.idle":"2022-11-18T18:32:35.218325Z","shell.execute_reply.started":"2022-11-18T18:32:35.209756Z","shell.execute_reply":"2022-11-18T18:32:35.217616Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### Dataset combinaton example","metadata":{"execution":{"iopub.status.busy":"2021-11-26T09:08:02.504247Z","iopub.execute_input":"2021-11-26T09:08:02.504566Z","iopub.status.idle":"2021-11-26T09:08:02.508303Z","shell.execute_reply.started":"2021-11-26T09:08:02.504533Z","shell.execute_reply":"2021-11-26T09:08:02.507579Z"}}},{"cell_type":"code","source":"# in general yo uwant to batch after shuffling but it doesn't matter for batch_size=1\nds_A = tf.data.Dataset.from_tensor_slices(['a', 'b', 'c']).batch(1).shuffle(3, reshuffle_each_iteration=True).repeat()\nds_B = tf.data.Dataset.range(5).batch(1).shuffle(5, reshuffle_each_iteration=True).repeat()\nds_zipped_dummy = tf.data.Dataset.zip((ds_A, ds_B))\n\nlist(ds_zipped_dummy.take(10).as_numpy_iterator())","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:35.219701Z","iopub.execute_input":"2022-11-18T18:32:35.220157Z","iopub.status.idle":"2022-11-18T18:32:35.244563Z","shell.execute_reply.started":"2022-11-18T18:32:35.220124Z","shell.execute_reply":"2022-11-18T18:32:35.243939Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Custom Callbacks","metadata":{}},{"cell_type":"code","source":"class ShowSamplesCallback(keras.callbacks.Callback):\n    def __init__(self, save_to_dir: str):\n        self.save_to_dir = save_to_dir\n    \n    def on_train_begin(self, logs=None):\n        self.original = []\n        self.styled = []\n    \n    def on_train_end(self, logs=None):\n        if self.save_to_dir:\n            dir_ = Path(self.save_to_dir)\n            dir_.mkdir(parents=True, exist_ok=True)\n            for i, (original, styled) in enumerate(zip(self.original, self.styled)):\n                tf.keras.preprocessing.image.save_img(dir_ / f'{i + 1}_original.jpg', original)\n                tf.keras.preprocessing.image.save_img(dir_ / f'{i + 1}_styled.jpg', styled)\n    \n    def on_epoch_end(self, epoch, logs=None):\n        styled_generator = self.model.s_gen\n        # sample predict\n        img_batch = next(iter(photo_ds.take(1)))\n        prediction = styled_generator(img_batch, training=False)[0].numpy()\n        prediction = (prediction * 127.5 + 127.5).astype(np.uint8)\n        img = img_batch[0]\n        img = (img * 127.5 + 127.5).numpy().astype(np.uint8)\n        self.original.append(img)\n        self.styled.append(prediction)\n        \n        # plot\n        fig, axs = plt.subplots(1, 2, figsize=(6,6), squeeze=True)\n        axs[0].imshow(img)\n        axs[1].imshow(prediction)\n        axs[0].set_title(\"Input Photo [{:02d}]\".format(epoch + 1))\n        axs[1].set_title(\"Generated [{:02d}]\".format(epoch + 1))\n        axs[0].axis(\"off\")\n        axs[1].axis(\"off\")\n        plt.show()","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:35.247126Z","iopub.execute_input":"2022-11-18T18:32:35.247322Z","iopub.status.idle":"2022-11-18T18:32:35.257201Z","shell.execute_reply.started":"2022-11-18T18:32:35.247300Z","shell.execute_reply":"2022-11-18T18:32:35.256496Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class SaveGeneratorCallback(keras.callbacks.Callback):\n    def __init__(self, epochs: int, path: str):\n        self.epochs = epochs\n        self.dir_ = Path(path)\n    \n    def on_epoch_end(self, epoch, logs=None):\n        if epoch % self.epochs == 0:\n            styled_generator = self.model.s_gen\n            path = self.dir_ / f'weights.{epoch + 1}.hdf5'\n            print(f'[SaveGeneratorCallback] Saving weights to {path}.')\n            styled_generator.save_weights(path)","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:35.258236Z","iopub.execute_input":"2022-11-18T18:32:35.259143Z","iopub.status.idle":"2022-11-18T18:32:35.272690Z","shell.execute_reply.started":"2022-11-18T18:32:35.259085Z","shell.execute_reply":"2022-11-18T18:32:35.271946Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Training","metadata":{}},{"cell_type":"code","source":"HOUR_TO_SECONDS = 3600","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:35.274194Z","iopub.execute_input":"2022-11-18T18:32:35.274467Z","iopub.status.idle":"2022-11-18T18:32:35.284769Z","shell.execute_reply.started":"2022-11-18T18:32:35.274426Z","shell.execute_reply":"2022-11-18T18:32:35.283936Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"_buffer_size = 3000\n\nhistory = cycle_gan_model.fit(\n    # zipped dataset - infinitely repeating\n    tf.data.Dataset.zip((\n        # buffer_size should be at least the size of the dataset for uniform shuffling\n        # - https://stackoverflow.com/a/47025850/3936732\n        # - however the dataset is too large for that so we limit the size\n        styled_ds.shuffle(_buffer_size, reshuffle_each_iteration=True).repeat(),\n        photo_ds.shuffle(_buffer_size, reshuffle_each_iteration=True).repeat(),\n    )),\n    \n    steps_per_epoch=1000, # batches per epoch\n    epochs=50,\n    # --- callbacks ---\n    callbacks=[\n        tfa.callbacks.TimeStopping(seconds=HOUR_TO_SECONDS * 6, verbose=1),\n        ShowSamplesCallback('/kaggle/working/imgs'),\n        SaveGeneratorCallback(50, '/kaggle/working'),\n              ],\n)","metadata":{"execution":{"iopub.status.busy":"2022-11-18T18:32:35.286903Z","iopub.execute_input":"2022-11-18T18:32:35.287164Z","iopub.status.idle":"2022-11-18T22:00:42.307615Z","shell.execute_reply.started":"2022-11-18T18:32:35.287133Z","shell.execute_reply":"2022-11-18T22:00:42.305779Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Visualize","metadata":{}},{"cell_type":"code","source":"_, ax = plt.subplots(2, 5, figsize=(16, 12))\nfor i, img in enumerate(photo_ds.take(5)):\n    prediction = styled_generator(img, training=False)[0].numpy()\n    prediction = (prediction * 127.5 + 127.5).astype(np.uint8)\n    img = (img[0] * 127.5 + 127.5).numpy().astype(np.uint8)\n\n    ax[0, i].imshow(img)\n    ax[1, i].imshow(prediction)\n    ax[0, i].set_title(\"Input Photo\")\n    ax[1, i].set_title(\"Generated\")\n    ax[0, i].axis(\"off\")\n    ax[1, i].axis(\"off\")\nplt.tight_layout()\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-11-18T22:00:42.312005Z","iopub.execute_input":"2022-11-18T22:00:42.312650Z","iopub.status.idle":"2022-11-18T22:00:43.430116Z","shell.execute_reply.started":"2022-11-18T22:00:42.312620Z","shell.execute_reply":"2022-11-18T22:00:43.429447Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model Saving/Loading","metadata":{}},{"cell_type":"markdown","source":"## Loading weights only (saved with callback)","metadata":{}},{"cell_type":"code","source":"_TEST = False","metadata":{"execution":{"iopub.status.busy":"2022-11-18T22:00:43.431280Z","iopub.execute_input":"2022-11-18T22:00:43.431636Z","iopub.status.idle":"2022-11-18T22:00:43.435956Z","shell.execute_reply.started":"2022-11-18T22:00:43.431606Z","shell.execute_reply":"2022-11-18T22:00:43.434944Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if _TEST:\n    generator = Generator()\n    generator.load_weights('/kaggle/working/weights.1.hdf5')\n    sample_photo = next(iter(photo_ds))\n    prediction = generator(sample_photo)\n    plot(sample_photo[0])\n    plt.show()\n    plot(prediction[0])\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2022-11-18T22:00:43.437338Z","iopub.execute_input":"2022-11-18T22:00:43.437660Z","iopub.status.idle":"2022-11-18T22:00:43.446417Z","shell.execute_reply.started":"2022-11-18T22:00:43.437622Z","shell.execute_reply":"2022-11-18T22:00:43.445607Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Final Model","metadata":{}},{"cell_type":"code","source":"!mkdir -p /kaggle/working/saved_model\nstyled_generator.save('/kaggle/working/saved_model/styled_generator')","metadata":{"execution":{"iopub.status.busy":"2022-11-18T22:00:43.465403Z","iopub.execute_input":"2022-11-18T22:00:43.465700Z","iopub.status.idle":"2022-11-18T22:00:54.899755Z","shell.execute_reply.started":"2022-11-18T22:00:43.465663Z","shell.execute_reply":"2022-11-18T22:00:54.898957Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"loaded_styled_generator = tf.keras.models.load_model('/kaggle/working/saved_model/styled_generator')","metadata":{"execution":{"iopub.status.busy":"2022-11-18T22:00:54.927248Z","iopub.execute_input":"2022-11-18T22:00:54.927525Z","iopub.status.idle":"2022-11-18T22:01:00.871969Z","shell.execute_reply.started":"2022-11-18T22:00:54.927490Z","shell.execute_reply":"2022-11-18T22:01:00.871175Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"n_rows = 3 * 2\nn_cols = 5\n_, ax = plt.subplots(n_rows, n_cols, figsize=(n_cols*4, n_rows*4))\nrow = 0\ni = 0\nfor img in photo_ds.take(n_rows//2*n_cols):\n    prediction = loaded_styled_generator(img, training=False)[0].numpy()\n    prediction = (prediction * 127.5 + 127.5).astype(np.uint8)\n    img = (img[0] * 127.5 + 127.5).numpy().astype(np.uint8)\n    \n    ax[row + 0, i].imshow(img)\n    ax[row + 1, i].imshow(prediction)\n    ax[row + 0, i].set_title(\"Input Photo\")\n    ax[row + 1, i].set_title(\"Generated\")\n    ax[row + 0, i].axis(\"off\")\n    ax[row + 1, i].axis(\"off\")\n    \n    i += 1\n    if i >= n_cols:\n        i = 0\n        row += 2\n    \nplt.tight_layout()\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-11-18T22:01:00.875819Z","iopub.execute_input":"2022-11-18T22:01:00.876062Z","iopub.status.idle":"2022-11-18T22:01:04.274405Z","shell.execute_reply.started":"2022-11-18T22:01:00.876036Z","shell.execute_reply":"2022-11-18T22:01:04.273331Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}