{"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":"import tensorflow as tf\n\ntpu = tf.distribute.cluster_resolver.TPUClusterResolver()\ntf.config.experimental_connect_to_cluster(tpu)\ntf.tpu.experimental.initialize_tpu_system(tpu)\nstrategy = tf.distribute.experimental.TPUStrategy(tpu)","metadata":{"execution":{"iopub.status.busy":"2023-03-01T08:57:15.453981Z","iopub.execute_input":"2023-03-01T08:57:15.455523Z","iopub.status.idle":"2023-03-01T08:57:27.602165Z","shell.execute_reply.started":"2023-03-01T08:57:15.455387Z","shell.execute_reply":"2023-03-01T08:57:27.601289Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from kaggle_datasets import KaggleDatasets\n\ngcs_path = KaggleDatasets().get_gcs_path()\ngcs_path","metadata":{"execution":{"iopub.status.busy":"2023-03-01T08:57:27.604765Z","iopub.execute_input":"2023-03-01T08:57:27.605862Z","iopub.status.idle":"2023-03-01T08:57:27.948126Z","shell.execute_reply.started":"2023-03-01T08:57:27.605791Z","shell.execute_reply":"2023-03-01T08:57:27.947072Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"AUTOTUNE = tf.data.experimental.AUTOTUNE\nIMAGE_SIZE = [256, 256]\n\nmonet_filenames = tf.io.gfile.glob(gcs_path + '/monet_tfrec/*.tfrec')\nphoto_filenames = tf.io.gfile.glob(gcs_path + '/photo_tfrec/*.tfrec')\n\n","metadata":{"execution":{"iopub.status.busy":"2023-03-01T08:57:27.949597Z","iopub.execute_input":"2023-03-01T08:57:27.950357Z","iopub.status.idle":"2023-03-01T08:57:28.050540Z","shell.execute_reply.started":"2023-03-01T08:57:27.950322Z","shell.execute_reply":"2023-03-01T08:57:28.049032Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"features = {'image': tf.io.FixedLenFeature([], tf.string)}\n\ndef read_tfrecord(example):\n    image_data = tf.io.parse_single_example(example, features)\n    image = image_data['image']\n    image = tf.image.decode_jpeg(image, channels=3)\n    image = (tf.cast(image, tf.float32) / 127.5) - 1\n    image = tf.reshape(image, [*IMAGE_SIZE, 3])\n    return image\n\n\nmonet_dataset = tf.data.TFRecordDataset(monet_filenames).map(read_tfrecord, num_parallel_calls=AUTOTUNE).batch(1)\nphoto_dataset = tf.data.TFRecordDataset(photo_filenames).map(read_tfrecord, num_parallel_calls=AUTOTUNE).batch(1)","metadata":{"execution":{"iopub.status.busy":"2023-03-01T08:57:28.054159Z","iopub.execute_input":"2023-03-01T08:57:28.054394Z","iopub.status.idle":"2023-03-01T08:57:28.164952Z","shell.execute_reply.started":"2023-03-01T08:57:28.054368Z","shell.execute_reply":"2023-03-01T08:57:28.163806Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"monet_dataset_len = len(list(iter(monet_dataset)))\nphoto_dataset_len = len(list(photo_dataset.as_numpy_iterator()))\nprint(monet_dataset_len, photo_dataset_len)","metadata":{"execution":{"iopub.status.busy":"2023-03-01T08:57:28.166331Z","iopub.execute_input":"2023-03-01T08:57:28.166659Z","iopub.status.idle":"2023-03-01T08:58:00.889573Z","shell.execute_reply.started":"2023-03-01T08:57:28.166618Z","shell.execute_reply":"2023-03-01T08:58:00.888617Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"example = next(iter(monet_dataset))","metadata":{"execution":{"iopub.status.busy":"2023-03-01T08:58:00.892545Z","iopub.execute_input":"2023-03-01T08:58:00.892728Z","iopub.status.idle":"2023-03-01T08:58:00.933359Z","shell.execute_reply.started":"2023-03-01T08:58:00.892706Z","shell.execute_reply":"2023-03-01T08:58:00.932232Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"example.shape","metadata":{"execution":{"iopub.status.busy":"2023-03-01T08:58:00.934710Z","iopub.execute_input":"2023-03-01T08:58:00.934982Z","iopub.status.idle":"2023-03-01T08:58:00.941566Z","shell.execute_reply.started":"2023-03-01T08:58:00.934947Z","shell.execute_reply":"2023-03-01T08:58:00.940567Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import tensorflow as tf\nimport tensorflow_addons as tfa\nfrom tensorflow import keras\nfrom tensorflow.keras import layers\n\nOUTPUT_CHANNELS = 3\n\n\nclass CycleGan(keras.Model):\n    def __init__(self):\n        super(CycleGan, self).__init__()\n        self.monet_generator = self.create_generator()\n        self.photo_generator = self.create_generator()\n        self.monet_discriminator = self.create_discriminator()\n        self.photo_discriminator = self.create_discriminator()\n        self.lambda_cycle = 10\n\n    def compile(self):\n        super(CycleGan, self).compile()\n        self.monet_generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n        self.monet_discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n        self.photo_generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n        self.photo_discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n\n    def create_downsampler(self, filters, size, apply_instance_norm=True):\n        model = keras.Sequential()\n        model.add(\n            layers.Conv2D(\n                filters,\n                size,\n                strides=2,\n                padding=\"same\",\n                use_bias=False,\n                kernel_initializer=tf.random_normal_initializer(0.0, 0.02),\n            )\n        )\n        if apply_instance_norm:\n            model.add(\n                tfa.layers.InstanceNormalization(\n                    gamma_initializer=keras.initializers.RandomNormal(\n                        mean=0.0, stddev=0.02\n                    )\n                )\n            )\n        model.add(layers.LeakyReLU())\n        return model\n\n    def create_upsampler(self, filters, size, apply_dropout=False):\n        model = keras.Sequential()\n        model.add(\n            layers.Conv2DTranspose(\n                filters,\n                size,\n                strides=2,\n                padding=\"same\",\n                use_bias=False,\n                kernel_initializer=tf.random_normal_initializer(0.0, 0.02),\n            )\n        )\n        model.add(\n            tfa.layers.InstanceNormalization(\n                gamma_initializer=tf.random_normal_initializer(0.0, 0.02)\n            )\n        )\n        if apply_dropout:\n            model.add(layers.Dropout(0.5))\n        model.add(layers.ReLU())\n        return model\n\n    def create_generator(self):\n        downsampler_stack = [\n            self.create_downsampler(64, 4, apply_instance_norm=False),\n            self.create_downsampler(128, 4),\n            self.create_downsampler(256, 4),\n        ] + [self.create_downsampler(512, 4) for i in range(5)]\n        upsampler_stack = [\n            self.create_upsampler(512, 4, apply_dropout=True) for i in range(3)\n        ] + [\n            self.create_upsampler(512, 4),\n            self.create_upsampler(256, 4),\n            self.create_upsampler(128, 4),\n            self.create_upsampler(64, 4),\n        ]\n        input_layer = layers.Input(shape=[256, 256, 3])\n        x = input_layer\n        skips = []\n        for downsampler in downsampler_stack:\n            x = downsampler(x)\n            skips.append(x)\n        skips = reversed(skips[:-1])\n        for upsampler, skip_layer in zip(upsampler_stack, skips):\n            x = upsampler(x)\n            x = layers.Concatenate()([x, skip_layer])\n        last_layer = layers.Conv2DTranspose(\n            OUTPUT_CHANNELS,\n            4,\n            strides=2,\n            padding=\"same\",\n            kernel_initializer=tf.random_normal_initializer(0.0, 0.02),\n            activation=\"tanh\",\n        )\n        x = last_layer(x)\n        return keras.Model(inputs=input_layer, outputs=x)\n\n    def create_discriminator(self):\n        input_layer = layers.Input(shape=[256, 256, 3], name=\"input_image\")\n        x = input_layer\n        downsampler1 = self.create_downsampler(64, 4, False)(x)\n        downsampler2 = self.create_downsampler(128, 4)(downsampler1)\n        downsampler3 = self.create_downsampler(256, 4)(downsampler2)\n        zero_pad1 = layers.ZeroPadding2D()(downsampler3)\n        conv_layer = layers.Conv2D(\n            512,\n            4,\n            strides=1,\n            use_bias=False,\n            kernel_initializer=tf.random_normal_initializer(0.0, 0.02),\n        )(zero_pad1)\n        normalization_layer1 = tfa.layers.InstanceNormalization(\n            gamma_initializer=tf.random_normal_initializer(0.0, 0.02)\n        )(conv_layer)\n        leaky_relu_layer = layers.LeakyReLU()(normalization_layer1)\n        zero_pad2 = layers.ZeroPadding2D()(leaky_relu_layer)\n        last_layer = layers.Conv2D(\n            1, 4, strides=1, kernel_initializer=tf.random_normal_initializer(0.0, 0.02)\n        )(zero_pad2)\n        return tf.keras.Model(inputs=input_layer, outputs=last_layer)\n\n    def discriminator_loss_fn(self, real, fake):\n        real_loss = tf.keras.losses.BinaryCrossentropy(\n            from_logits=True, reduction=tf.keras.losses.Reduction.NONE\n        )(tf.ones_like(real), real)\n        fake_loss = tf.keras.losses.BinaryCrossentropy(\n            from_logits=True, reduction=tf.keras.losses.Reduction.NONE\n        )(tf.zeros_like(fake), fake)\n        return (real_loss + fake_loss) / 2\n\n    def generator_loss_fn(self, generated_image):\n        return tf.keras.losses.BinaryCrossentropy(\n            from_logits=True, reduction=tf.keras.losses.Reduction.NONE\n        )(tf.ones_like(generated_image), generated_image)\n\n    def cycle_loss_fn(self, image, cycled_image, lambda_cycle):\n        return tf.reduce_mean(tf.abs(image - cycled_image)) * lambda_cycle\n\n    def identity_loss_fn(self, real_photo, photo, lambda_cycle):\n        return tf.reduce_mean(tf.abs(real_photo - photo)) * lambda_cycle / 2\n\n    def train_step(self, batch_data):\n        real_monet, real_photo = batch_data\n\n        with tf.GradientTape(persistent=True) as tape:\n            fake_monet = self.monet_generator(real_photo, training=True)\n            cycled_photo = self.photo_generator(fake_monet, training=True)\n            fake_photo = self.photo_generator(real_monet, training=True)\n            cycled_monet = self.monet_generator(fake_photo, training=True)\n\n            monet1 = self.monet_generator(real_monet, training=True)\n            photo1 = self.photo_generator(real_photo, training=True)\n\n            monet_real_discriminated = self.monet_discriminator(\n                real_monet, training=True\n            )\n            monet_fake_discriminated = self.monet_discriminator(\n                fake_monet, training=True\n            )\n            photo_real_discriminated = self.photo_discriminator(\n                real_photo, training=True\n            )\n            photo_fake_discriminated = self.photo_discriminator(\n                fake_photo, training=True\n            )\n\n            monet_generator_loss = self.generator_loss_fn(monet_fake_discriminated)\n            photo_generator_loss = self.generator_loss_fn(photo_fake_discriminated)\n            cycle_loss = self.cycle_loss_fn(\n                real_monet, cycled_monet, self.lambda_cycle\n            ) + self.cycle_loss_fn(real_photo, cycled_photo, self.lambda_cycle)\n            total_monet_generator_loss = (\n                monet_generator_loss\n                + cycle_loss\n                + self.identity_loss_fn(real_monet, monet1, self.lambda_cycle)\n            )\n            total_photo_generator_loss = (\n                photo_generator_loss\n                + cycle_loss\n                + self.identity_loss_fn(real_photo, photo1, self.lambda_cycle)\n            )\n            monet_discriminator_loss = self.discriminator_loss_fn(\n                monet_real_discriminated, monet_fake_discriminated\n            )\n            photo_discriminator_loss = self.discriminator_loss_fn(\n                photo_real_discriminated, photo_fake_discriminated\n            )\n        monet_generator_gradients = tape.gradient(\n            total_monet_generator_loss, self.monet_generator.trainable_variables\n        )\n        photo_generator_gradients = tape.gradient(\n            total_photo_generator_loss, self.photo_generator.trainable_variables\n        )\n        monet_discriminator_gradients = tape.gradient(\n            monet_discriminator_loss, self.monet_discriminator.trainable_variables\n        )\n        photo_discriminator_gradients = tape.gradient(\n            photo_discriminator_loss, self.photo_discriminator.trainable_variables\n        )\n        self.monet_generator_optimizer.apply_gradients(\n            zip(monet_generator_gradients, self.monet_generator.trainable_variables)\n        )\n        self.photo_generator_optimizer.apply_gradients(\n            zip(photo_generator_gradients, self.photo_generator.trainable_variables)\n        )\n        self.monet_discriminator_optimizer.apply_gradients(\n            zip(\n                monet_discriminator_gradients,\n                self.monet_discriminator.trainable_variables,\n            )\n        )\n        self.photo_discriminator_optimizer.apply_gradients(\n            zip(\n                photo_discriminator_gradients,\n                self.photo_discriminator.trainable_variables,\n            )\n        )\n        return {\n            \"monet_generator_loss\": total_monet_generator_loss,\n            \"photo_generator_loss\": total_photo_generator_loss,\n            \"monet_discriminator_loss\": monet_discriminator_loss,\n            \"photo_discriminator_loss\": photo_discriminator_loss,\n        }","metadata":{"_kg_hide-input":false,"execution":{"iopub.status.busy":"2023-03-01T08:58:00.943108Z","iopub.execute_input":"2023-03-01T08:58:00.943432Z","iopub.status.idle":"2023-03-01T08:58:01.137074Z","shell.execute_reply.started":"2023-03-01T08:58:00.943382Z","shell.execute_reply":"2023-03-01T08:58:01.135925Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with strategy.scope():\n    cycle_gan_model = CycleGan()\n    cycle_gan_model.compile()","metadata":{"execution":{"iopub.status.busy":"2023-03-01T08:58:01.138638Z","iopub.execute_input":"2023-03-01T08:58:01.138898Z","iopub.status.idle":"2023-03-01T08:58:10.419522Z","shell.execute_reply.started":"2023-03-01T08:58:01.138869Z","shell.execute_reply":"2023-03-01T08:58:10.418617Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cycle_gan_model.fit(\n    tf.data.Dataset.zip((monet_dataset.repeat(-1), photo_dataset.repeat(-1))),\n    steps_per_epoch=300,\n    epochs=30\n)","metadata":{"execution":{"iopub.status.busy":"2023-03-01T08:58:10.421923Z","iopub.execute_input":"2023-03-01T08:58:10.422183Z","iopub.status.idle":"2023-03-01T09:19:03.143115Z","shell.execute_reply.started":"2023-03-01T08:58:10.422155Z","shell.execute_reply":"2023-03-01T09:19:03.141798Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from io import BytesIO\nfrom PIL import Image\nfrom zipfile import ZipFile\nimport numpy as np\n\nwith ZipFile('images.zip', mode='w') as zip_file:\n    i = 1\n    for img in photo_dataset:\n        generated_image_data = cycle_gan_model.monet_generator(img, training=False)[0].numpy()\n        scaled_generated_image_data = (generated_image_data * 127.5 + 127.5).astype(np.uint8)\n        with BytesIO() as image_bytes_io:\n            Image.fromarray(scaled_generated_image_data).save(image_bytes_io, 'JPEG')\n            image_bytes_io.seek(0)\n            zip_file.writestr('{}.jpg'.format(i), image_bytes_io.read())\n            i += 1","metadata":{"execution":{"iopub.status.busy":"2023-03-01T09:24:17.887655Z","iopub.execute_input":"2023-03-01T09:24:17.888009Z","iopub.status.idle":"2023-03-01T10:01:31.047903Z","shell.execute_reply.started":"2023-03-01T09:24:17.887975Z","shell.execute_reply":"2023-03-01T10:01:31.046609Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}