{"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 os\nfrom glob import glob\nimport zipfile\nimport shutil\nimport time\n\nimport numpy as np\nimport tensorflow as tf\nimport matplotlib.pyplot as plt\nfrom IPython.display import clear_output\n\nAUTOTUNE = tf.data.AUTOTUNE","metadata":{"_uuid":"ea727fff-0912-4cda-ab6d-d34f4b8682d6","_cell_guid":"aeebbb30-1051-4cc1-bb33-3ba822b1b2c8","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-10-03T22:08:16.660194Z","iopub.execute_input":"2022-10-03T22:08:16.660764Z","iopub.status.idle":"2022-10-03T22:08:21.157093Z","shell.execute_reply.started":"2022-10-03T22:08:16.660644Z","shell.execute_reply":"2022-10-03T22:08:21.156112Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Comenzamos extrayendo las imágenes del conjunto de entrenamiento y prueba para almacenarlas en las carpetas `./train/` y `./test1/`.","metadata":{"_uuid":"2eebf138-3b59-4b04-ad00-fe250fe1386b","_cell_guid":"cb9117cf-d514-498a-b2cb-27bda211ceb4","trusted":true}},{"cell_type":"code","source":"base_dir = \"../input/dogs-vs-cats\"\ntrain_dir = os.path.join(base_dir, \"train.zip\")\ntest_dir = os.path.join(base_dir, \"test1.zip\")\n\nwith zipfile.ZipFile(train_dir,\"r\") as z:\n    z.extractall()\n\nwith zipfile.ZipFile(test_dir,\"r\") as z:\n    z.extractall()","metadata":{"_uuid":"4be1a5eb-119c-4a4e-bacf-fb02c41a2c71","_cell_guid":"9d3677c7-66f2-40a2-9e80-712a45a1ae2c","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-10-03T22:08:21.159206Z","iopub.execute_input":"2022-10-03T22:08:21.159879Z","iopub.status.idle":"2022-10-03T22:08:39.973815Z","shell.execute_reply.started":"2022-10-03T22:08:21.159817Z","shell.execute_reply":"2022-10-03T22:08:39.972822Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Luego, imprimimos la cantidad de imágenes que existen por cada categoría.","metadata":{"_uuid":"9499ca6b-9c22-48f2-835d-3a428490e1f2","_cell_guid":"e26252a4-2e24-4886-8364-dce2ff3aba18","trusted":true}},{"cell_type":"code","source":"print(f\"Number of dogs: {len(glob('./train/dog*'))}\\nNumber of cats: {len(glob('./train/cat*'))}\")","metadata":{"_uuid":"b13b62d1-ea6c-4275-a4bf-b48e2cac9fa9","_cell_guid":"0e89482a-3b34-40dd-bb46-fa21ecad8396","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-10-03T22:08:39.975383Z","iopub.execute_input":"2022-10-03T22:08:39.975756Z","iopub.status.idle":"2022-10-03T22:08:40.076864Z","shell.execute_reply.started":"2022-10-03T22:08:39.975717Z","shell.execute_reply":"2022-10-03T22:08:40.075856Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Si queremos utilizar la función de Keras `image_dataset_from_directory`, movemos las imágenes de perros a la carpeta `./train/dog/` y las imágenes de gatos a la carpeta `./train/cat/`.","metadata":{}},{"cell_type":"code","source":"!mkdir ./train/dog/\n!mkdir ./train/cat/\n!mv ./train/dog.* ./train/dog/\n!mv ./train/cat.* ./train/cat/","metadata":{"execution":{"iopub.status.busy":"2022-10-03T22:08:40.079581Z","iopub.execute_input":"2022-10-03T22:08:40.080349Z","iopub.status.idle":"2022-10-03T22:08:45.014894Z","shell.execute_reply.started":"2022-10-03T22:08:40.080319Z","shell.execute_reply":"2022-10-03T22:08:45.013458Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Una vez organizados nuestros archivos, definimos un conjunto de variables globales que serán utilizadas a lo largo de todo el código.","metadata":{"_uuid":"c2655aa8-4dc3-4bc6-af21-182d0cb10de9","_cell_guid":"ea42b3cb-7c9b-4bf0-9e21-8fbf699e6859","trusted":true}},{"cell_type":"code","source":"BATCH_SIZE = 16\nIMG_WIDTH = 256\nIMG_HEIGHT = 256\nOUTPUT_CHANNELS = 3\n\n# Normalization layer to represent images in [-1,1] interval\nnorm_layer = tf.keras.layers.experimental.preprocessing.Rescaling(scale=1./127.5, offset=-1)","metadata":{"_uuid":"ee66760f-59ef-4d35-9a01-54ab5a0e49d9","_cell_guid":"4b065c02-3674-484a-b938-6f56137909bf","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-10-03T22:08:45.017960Z","iopub.execute_input":"2022-10-03T22:08:45.018613Z","iopub.status.idle":"2022-10-03T22:08:45.937614Z","shell.execute_reply.started":"2022-10-03T22:08:45.018568Z","shell.execute_reply":"2022-10-03T22:08:45.936658Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Ahora, para acelerar un poco el entrenamiento a costa de algo de pérdida de la capacidad de generalización de nuestro modelo, seleccionamos`BATCH_SIZE*100` imágenes de las 12500 imágenes por categoría. En otras palabras, utilizaremos únicamente `BATCH_SIZE*100` imágenes de perros y `BATCH_SIZE*100` imágenes de gatos. ","metadata":{"_uuid":"a9caa9b7-a6fb-4b32-aa1a-382c66f91972","_cell_guid":"15c5f964-cd34-4667-b989-ad4ebcfee7af","execution":{"iopub.status.busy":"2022-09-27T22:43:46.164898Z","iopub.execute_input":"2022-09-27T22:43:46.165342Z","iopub.status.idle":"2022-09-27T22:43:46.172424Z","shell.execute_reply.started":"2022-09-27T22:43:46.165303Z","shell.execute_reply":"2022-09-27T22:43:46.171018Z"},"trusted":true}},{"cell_type":"code","source":"train_X = tf.keras.utils.image_dataset_from_directory(directory='./train/dog/', labels=None, seed=123, batch_size=BATCH_SIZE).take(100)\ntrain_X = train_X.map(lambda image: norm_layer(image), num_parallel_calls=AUTOTUNE)\n\ntrain_Y = tf.keras.utils.image_dataset_from_directory(directory='./train/cat/', labels=None, seed=123, batch_size=BATCH_SIZE).take(100)\ntrain_Y = train_Y.map(lambda image: norm_layer(image), num_parallel_calls=AUTOTUNE)\n\ntrain = tf.data.Dataset.zip((train_X, train_Y))","metadata":{"_uuid":"a365e324-a55d-4db8-82e9-4ef084f15c48","_cell_guid":"f10865f0-ab70-411e-af9a-b06b9b458aed","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-10-03T22:08:45.939321Z","iopub.execute_input":"2022-10-03T22:08:45.939683Z","iopub.status.idle":"2022-10-03T22:08:49.899794Z","shell.execute_reply.started":"2022-10-03T22:08:45.939647Z","shell.execute_reply":"2022-10-03T22:08:49.898699Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Definición del modelo Generador\nEl modelo generador consta de un autoencoder variacional (VAE), que recibe una entrada  cuya salida se conecta a la entrada de una red tipo UNet.","metadata":{"_uuid":"bfb6dbd1-d3ba-4445-b432-73a102859a2d","_cell_guid":"53dd967f-f06e-405a-9f5b-e22c56c8bd20","trusted":true}},{"cell_type":"code","source":"# Instance Normalization Layer\nclass InstanceNormalization(tf.keras.layers.Layer):\n    \"\"\"Instance Normalization Layer (https://arxiv.org/abs/1607.08022).\"\"\"\n\n    def __init__(self, epsilon=1e-5):\n        super(InstanceNormalization, self).__init__()\n        self.epsilon = epsilon\n\n    def build(self, input_shape):\n        self.scale = self.add_weight(name='scale',\n                                     shape=input_shape[-1:],\n                                     initializer=tf.random_normal_initializer(1., 0.02),\n                                     trainable=True)\n\n        self.offset = self.add_weight(name='offset',\n                                      shape=input_shape[-1:], \n                                      initializer='zeros', \n                                      trainable=True)\n\n    def call(self, x):\n        mean, variance = tf.nn.moments(x, axes=[1, 2], keepdims=True)\n        inv = tf.math.rsqrt(variance + self.epsilon)\n        normalized = (x - mean) * inv\n        return self.scale * normalized + self.offset\n\n# Down-sampling block\ndef downsample(filters, size, apply_batchnorm=True):\n    initializer = tf.random_normal_initializer(0., 0.02)\n\n    result = tf.keras.Sequential()\n    result.add(tf.keras.layers.Conv2D(filters=filters,\n                                      kernel_size=size,\n                                      strides=2,\n                                      padding='same',\n                                      kernel_initializer=initializer,\n                                      use_bias=False))\n    \n    if apply_batchnorm:\n        result.add(InstanceNormalization()) # Perform instance normalization\n\n    result.add(tf.keras.layers.LeakyReLU())\n\n    return result\n\n# Downsampling code verification\nsample_dog = next(iter(train_X))\nsample_cat = next(iter(train_Y))\ndown_model = downsample(3,4)\ndown_result = down_model(tf.expand_dims(sample_dog[0], 0))\n\n# Up-sampling block\ndef upsample(filters, size, apply_dropout=False):\n    initializer = tf.random_normal_initializer(0., 0.02)\n    \n    result = tf.keras.Sequential()\n    result.add(tf.keras.layers.Conv2DTranspose(filters=filters,\n                                               kernel_size=size,\n                                               strides=2,\n                                               padding='same',\n                                               kernel_initializer=initializer,\n                                               use_bias=False))\n\n    result.add(InstanceNormalization()) # Perform instance normalization\n\n    if apply_dropout:\n        result.add(tf.keras.layers.Dropout(0.5))\n\n    result.add(tf.keras.layers.ReLU())\n\n    return result\n\n# Upsampling verification\nup_model = upsample(3, 4)\nup_result = up_model(down_result)\nprint (f'Output dimensions{up_result.shape}')","metadata":{"_uuid":"136d5d6f-f039-4942-a927-d086c902ad45","_cell_guid":"100626fa-8b8e-47b3-9094-abe922178348","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-10-03T22:08:49.903136Z","iopub.execute_input":"2022-10-03T22:08:49.903423Z","iopub.status.idle":"2022-10-03T22:08:57.245864Z","shell.execute_reply.started":"2022-10-03T22:08:49.903396Z","shell.execute_reply":"2022-10-03T22:08:57.244692Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"'''class Reparameterize(tf.keras.layers.Layer):\n    def __init__(self, output_dim, **kwargs):    \n        self.output_dim = output_dim \n        super(Reparameterize, self).__init__(**kwargs)\n        \n    def call(self, inputs):\n        mu, log_sigma2 = args\n        epsilon = tf.keras.backend.random_normal(shape=(BATCH_SIZE, latent_dim))\n        return mu + tf.keras.backend.exp(log_sigma2 / 2) * epsilon'''\n\nclass CVAE(tf.keras.Model):\n    \"\"\"Convolutional variational autoencoder.\"\"\"\n\n    def __init__(self, latent_dim):\n        super(CVAE, self).__init__()\n        self.latent_dim = latent_dim\n        self.encoder = tf.keras.Sequential(\n            [\n                tf.keras.layers.InputLayer(input_shape=(256, 256, 3)),\n                tf.keras.layers.Conv2D(\n                    filters=32, kernel_size=3, strides=(2, 2), activation='relu'),\n                tf.keras.layers.Conv2D(\n                    filters=64, kernel_size=3, strides=(2, 2), activation='relu'),\n                tf.keras.layers.Flatten(),\n                # No activation\n                tf.keras.layers.Dense(latent_dim),\n            ]\n        )\n\n        self.decoder = tf.keras.Sequential(\n            [\n                tf.keras.layers.InputLayer(input_shape=(latent_dim,)),\n                tf.keras.layers.Dense(units=7*7*32, activation=tf.nn.relu),\n                tf.keras.layers.Reshape(target_shape=(7, 7, 32)),\n                tf.keras.layers.Conv2DTranspose(\n                    filters=64, kernel_size=3, strides=2, padding='same',\n                    activation='relu'),\n                tf.keras.layers.Conv2DTranspose(\n                    filters=32, kernel_size=3, strides=2, padding='same',\n                    activation='relu'),\n                # No activation\n                tf.keras.layers.Conv2DTranspose(\n                    filters=1, kernel_size=3, strides=1, padding='same'),\n        ]\n    )\n\n    @tf.function\n    def sample(self, eps=None):\n        if eps is None:\n            eps = tf.random.normal(shape=(100, self.latent_dim))\n        return self.decode(eps, apply_sigmoid=True)\n\n    def encode(self, x):\n        mean, logvar = tf.split(self.encoder(x), num_or_size_splits=2, axis=1)\n        return mean, logvar\n\n    def reparameterize(self, mean, logvar):\n        eps = tf.random.normal(shape=mean.shape)\n        return eps * tf.exp(logvar * .5) + mean\n\n    def decode(self, z, apply_sigmoid=False):\n        logits = self.decoder(z)\n        if apply_sigmoid:\n            probs = tf.sigmoid(logits)\n            return probs\n        return logits","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# VAE+UNet Generator\ndef Generator(latent_dim=3):\n    # Layers that compose the net\n    x_input = tf.keras.layers.Input(shape=(256, 256, 3))\n    \n    # Encoder \n    '''encoder_H1 = tf.keras.layers.Conv2D(filters=32, kernel_size=3, strides=(2, 2), activation='relu')(x_input)\n    encoder_H2 = tf.keras.layers.Conv2D(filters=64, kernel_size=3, strides=(2, 2), activation='relu')(encoder_H1)\n    flat = tf.keras.layers.Flatten()(encoder_H2)\n    \n    # Mean and log-variance of latent space\n    z_mu = tf.keras.layers.Dense(units=latent_dim, name=\"mean\")(flat)\n    z_log_sigma2 = tf.keras.layers.Dense(units=latent_dim, name=\"log-variance\")(flat)\n    \n    # Sampling z with the reparametrization trick\n    #z = tf.keras.layers.Lambda(sample_z, arguments=[z_mu, z_log_sigma2], output_shape=(latent_dim,))\n    \n    encoder = tf.keras.Model(inputs=x_input, outputs=[z_mu, z_log_sigma2], name=\"encoder\")\n    encoder.summary()\n    #tf.keras.utils.plot_model(encoder,show_shapes=True, show_layer_names=True)'''\n    \n    down_stack = [\n        downsample(64, 4, apply_batchnorm=False),  # (batch_size, 128, 128, 64)\n        downsample(128, 4),  # (batch_size, 64, 64, 128)\n        downsample(256, 4),  # (batch_size, 32, 32, 256)\n        downsample(512, 4),  # (batch_size, 16, 16, 512)\n        downsample(512, 4),  # (batch_size, 8, 8, 512)\n        downsample(512, 4),  # (batch_size, 4, 4, 512)\n        downsample(512, 4),  # (batch_size, 2, 2, 512)\n        downsample(512, 4),  # (batch_size, 1, 1, 512)\n    ]\n\n    up_stack = [\n        upsample(512, 4, apply_dropout=True),  # (batch_size, 2, 2, 1024)\n        upsample(512, 4, apply_dropout=True),  # (batch_size, 4, 4, 1024)\n        upsample(512, 4, apply_dropout=True),  # (batch_size, 8, 8, 1024)\n        upsample(512, 4),  # (batch_size, 16, 16, 1024)\n        upsample(256, 4),  # (batch_size, 32, 32, 512)\n        upsample(128, 4),  # (batch_size, 64, 64, 256)\n        upsample(64, 4),  # (batch_size, 128, 128, 128)\n    ]\n\n    initializer = tf.random_normal_initializer(0., 0.02)\n    last = tf.keras.layers.Conv2DTranspose(OUTPUT_CHANNELS,\n                                           4,\n                                           strides=2,\n                                           padding='same',\n                                           kernel_initializer=initializer,\n                                           activation='tanh') # (batch_size, 256, 256, 1)\n\n    # Processing pipeline\n    x = x_input\n    \n    # Encoder\n    skips=[]\n    for down in down_stack:\n        x = down(x)\n        skips.append(x) # Output for each downsampling is added to a list\n\n    skips = reversed(skips[:-1])\n    \n    #Decoder\n    for up, skip in zip(up_stack, skips):\n        x = up(x)\n        x = tf.keras.layers.Concatenate()([x, skip])\n    \n    x = last(x)\n\n    return tf.keras.Model(inputs=x_input, outputs=x)","metadata":{"_uuid":"b5290fc6-48c3-4bda-a2fb-e7f61d7a9277","_cell_guid":"441a8670-3c48-451b-8012-a208cfabbdda","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-10-03T22:08:57.258470Z","iopub.execute_input":"2022-10-03T22:08:57.258799Z","iopub.status.idle":"2022-10-03T22:08:57.270896Z","shell.execute_reply.started":"2022-10-03T22:08:57.258769Z","shell.execute_reply":"2022-10-03T22:08:57.269902Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Probamos las salidas de los generadores $G$ y $F$, con pesos inicializados aleatoriamente.","metadata":{"_uuid":"45d9813f-3d3e-4ec5-8312-4e60ed3eb007","_cell_guid":"2be07248-9de9-4f07-854e-352277dcf9cc","trusted":true}},{"cell_type":"code","source":"# Test generators\ngenerator_G = Generator()\ngen_output = generator_G(sample_dog[0][tf.newaxis, ...], training=False)\ngenerator_F = Generator()\ngen_output2 = generator_F(sample_cat[0][tf.newaxis, ...], training=False)\n\n'''plt.figure(figsize=(12,6))\nplt.subplot(121)\nplt.title('$G(X)$')\nplt.imshow(gen_output[0, ...]*50, cmap='gray')\nplt.axis('off')\nplt.subplot(122)\nplt.title('$F(Y)$')\nplt.imshow(gen_output2[0, ...]*50, cmap='gray')\nplt.axis('off')\nplt.show()'''","metadata":{"_uuid":"91ed009e-ed7a-4c04-b6c8-c96cf602a6ec","_cell_guid":"5a67c006-59c5-44f7-9634-72356004372d","collapsed":false,"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Ahora procedemos a definir nuestros discriminadores.","metadata":{"_uuid":"bd1b5a3a-6d4f-478d-88b3-c83edf3eca15","_cell_guid":"8460449c-85be-4668-acc7-4b188ca4e31e","trusted":true}},{"cell_type":"code","source":"def Discriminator():\n    initializer = tf.random_normal_initializer(0., 0.02)\n\n    x = tf.keras.layers.Input(shape=[256, 256, 3], name='input_image')\n\n    down1 = downsample(64, 4, False)(x)  # (batch_size, 128, 128, 64)\n    down2 = downsample(128, 4)(down1)  # (batch_size, 64, 64, 128)\n    down3 = downsample(256, 4)(down2)  # (batch_size, 32, 32, 256)\n\n    zero_pad1 = tf.keras.layers.ZeroPadding2D()(down3)  # (batch_size, 34, 34, 256)\n    conv = tf.keras.layers.Conv2D(512, 4, strides=1,\n                                    kernel_initializer=initializer,\n                                    use_bias=False)(zero_pad1)  # (batch_size, 31, 31, 512)\n\n    instnorm1 = InstanceNormalization()(conv)\n\n    leaky_relu = tf.keras.layers.LeakyReLU()(instnorm1)\n\n    zero_pad2 = tf.keras.layers.ZeroPadding2D()(leaky_relu)  # (batch_size, 33, 33, 512)\n\n    last = tf.keras.layers.Conv2D(1, 4, strides=1,\n                                    kernel_initializer=initializer)(zero_pad2)  # (batch_size, 30, 30, 1)\n\n    return tf.keras.Model(inputs=x, outputs=last)","metadata":{"_uuid":"cf14b6b0-d269-4c96-9846-d011dc91c3f0","_cell_guid":"16ed8d27-d267-4f60-b6fd-708a86116c9f","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-10-03T21:50:28.766679Z","iopub.status.idle":"2022-10-03T21:50:28.767405Z","shell.execute_reply.started":"2022-10-03T21:50:28.767173Z","shell.execute_reply":"2022-10-03T21:50:28.767196Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Test discriminators\ndiscriminator_X = Discriminator()\ndiscriminator_Y = Discriminator()\n\nplt.figure(figsize=(12, 6))\n\nplt.subplot(121)\nplt.title('Is a real dog?')\nplt.axis('off')\nplt.imshow(discriminator_X(sample_cat)[0, ..., -1], cmap='RdBu_r')\n\nplt.subplot(122)\nplt.title('Is a real cat?')\nplt.axis('off')\nplt.imshow(discriminator_Y(sample_dog)[0, ..., -1], cmap='RdBu_r')\n\nplt.show()","metadata":{"_uuid":"fa37a388-f27b-4933-b051-c5f6fe722def","_cell_guid":"bbeab99c-32bd-4b3f-9c62-afc6bfbe9fbd","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-10-03T21:50:28.775396Z","iopub.status.idle":"2022-10-03T21:50:28.776126Z","shell.execute_reply.started":"2022-10-03T21:50:28.775894Z","shell.execute_reply":"2022-10-03T21:50:28.775917Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Next, we define our loss functions.","metadata":{"_uuid":"b00f626e-34a7-40b0-9e16-1d1a5f654a25","_cell_guid":"c9b3b3cf-2076-4de5-a31f-352ba563bd2d","trusted":true}},{"cell_type":"code","source":"LAMBDA = 5\nloss_obj = tf.keras.losses.BinaryCrossentropy(from_logits=True)\n\ndef log_normal_pdf(sample, mean, logvar, raxis=1):\n    log2pi = tf.math.log(2. * np.pi)\n    return tf.reduce_sum(-.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi), axis=raxis)\n\n\ndef compute_loss(model, x):\n    mean, logvar = model.encode(x)\n    z = model.reparameterize(mean, logvar)\n    x_logit = model.decode(z)\n    cross_ent = tf.nn.sigmoid_cross_entropy_with_logits(logits=x_logit, labels=x)\n    logpx_z = -tf.reduce_sum(cross_ent, axis=[1, 2, 3])\n    logpz = log_normal_pdf(z, 0., 0.)\n    logqz_x = log_normal_pdf(z, mean, logvar)\n    return -tf.reduce_mean(logpx_z + logpz - logqz_x)\n\n\ndef discriminator_loss(real, generated):\n    real_loss = loss_obj(tf.ones_like(real), real)\n\n    generated_loss = loss_obj(tf.zeros_like(generated), generated)\n\n    total_disc_loss = real_loss + generated_loss\n\n    return total_disc_loss * 0.5\n\ndef generator_loss(discriminated):\n    return loss_obj(tf.ones_like(discriminated), discriminated)\n\ndef calc_cycle_loss(real_image, cycled_image):\n    loss1 = tf.reduce_mean(tf.abs(real_image - cycled_image))\n    \n    return LAMBDA * loss1\n\ndef identity_loss(real_image, same_image):\n    loss = tf.reduce_mean(tf.abs(real_image - same_image))\n    \n    return LAMBDA * 0.5 * loss","metadata":{"_uuid":"be9501aa-d77c-4b2a-ad13-722f2754afb9","_cell_guid":"538b3670-dc49-458f-9ee9-adda97bad97e","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-10-03T21:50:28.777412Z","iopub.status.idle":"2022-10-03T21:50:28.778140Z","shell.execute_reply.started":"2022-10-03T21:50:28.777900Z","shell.execute_reply":"2022-10-03T21:50:28.777923Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"generator_G_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\ngenerator_F_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n\ndiscriminator_X_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\ndiscriminator_Y_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)","metadata":{"_uuid":"0f498dfb-a943-4dc2-b7a2-94d85dc28152","_cell_guid":"27723086-ef94-4b02-86f3-1c648408c4ad","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-10-03T21:50:28.779412Z","iopub.status.idle":"2022-10-03T21:50:28.780179Z","shell.execute_reply.started":"2022-10-03T21:50:28.779930Z","shell.execute_reply":"2022-10-03T21:50:28.779953Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"checkpoint_path = \"./checkpoints/train\"\n\nckpt = tf.train.Checkpoint(generator_G=generator_G,\n                           generator_F=generator_F,\n                           discriminator_X=discriminator_X,\n                           discriminator_Y=discriminator_Y,\n                           generator_G_optimizer=generator_G_optimizer,\n                           generator_F_optimizer=generator_F_optimizer,\n                           discriminator_X_optimizer=discriminator_X_optimizer,\n                           discriminator_Y_optimizer=discriminator_Y_optimizer)\n\nckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=5)\n\n# if a checkpoint exists, restore the latest checkpoint.\nif ckpt_manager.latest_checkpoint:\n    ckpt.restore(ckpt_manager.latest_checkpoint)\n    print ('Latest checkpoint restored!!')","metadata":{"_uuid":"514bf793-8ed7-4771-91f6-3a2d0977524b","_cell_guid":"ffb3a581-943b-48a0-aa25-6b8588414048","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-10-03T21:50:28.781503Z","iopub.status.idle":"2022-10-03T21:50:28.782229Z","shell.execute_reply.started":"2022-10-03T21:50:28.781997Z","shell.execute_reply":"2022-10-03T21:50:28.782020Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training","metadata":{"_uuid":"09fa7ce1-7583-4226-8a26-0cfa69a71f73","_cell_guid":"bf57a72e-7671-49bd-bafb-adfca0254d23","trusted":true}},{"cell_type":"code","source":"def generate_images(model, test_input):\n    prediction = model(test_input)\n    \n    plt.figure(figsize=(12,12))\n    \n    display_list = [test_input[0], prediction[0]]\n    title = ['Input Image', 'Predicted Image']\n    \n    for i in range(2):\n        plt.subplot(1, 2, i+1)\n        plt.title(title[i])\n        # getting the pixel values between [0, 1] to plot it.\n        plt.imshow(display_list[i] * 0.5 + 0.5)\n        plt.axis('off')\n    plt.show()","metadata":{"_uuid":"29907fd5-79f6-42c8-8e10-02a32f163a9b","_cell_guid":"6b1436cc-1617-4d62-924c-0b52603ef144","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-10-03T21:50:28.783547Z","iopub.status.idle":"2022-10-03T21:50:28.784268Z","shell.execute_reply.started":"2022-10-03T21:50:28.784033Z","shell.execute_reply":"2022-10-03T21:50:28.784056Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"@tf.function\ndef train_step(real_X, real_Y):\n    # persistent is set to True because the tape is used more than\n    # once to calculate the gradients.\n    with tf.GradientTape(persistent=True) as tape:\n        # Generator G translates X -> Y\n        # Generator F translates Y -> X\n        \n        fake_Y = generator_G(real_X, training=True)\n        cycled_X = generator_F(fake_Y, training=True)\n        \n        fake_X = generator_F(real_Y, training=True)\n        cycled_Y = generator_G(fake_X, training=True)\n        \n        # same_X and same_Y are used for identity loss\n        same_X = generator_F(real_X, training=True)\n        same_Y = generator_G(real_Y, training=True)\n        \n        disc_real_X = discriminator_X(real_X, training=True)\n        disc_real_Y = discriminator_Y(real_Y, training=True)\n        \n        disc_fake_X = discriminator_X(fake_X, training=True)\n        disc_fake_Y = discriminator_Y(fake_Y, training=True)\n        \n        # calculate the loss\n        gen_G_loss = generator_loss(disc_fake_Y)\n        gen_F_loss = generator_loss(disc_fake_X)\n        \n        total_cycle_loss = calc_cycle_loss(real_X, cycled_X) + calc_cycle_loss(real_Y, cycled_Y)\n        \n        # total generator loss = adversarial loss + cycle loss\n        total_gen_G_loss = gen_G_loss + total_cycle_loss + identity_loss(real_Y, same_Y) \n        total_gen_F_loss = gen_F_loss + total_cycle_loss + identity_loss(real_X, same_X)\n        \n        disc_X_loss = discriminator_loss(disc_real_X, disc_fake_X)\n        disc_Y_loss = discriminator_loss(disc_real_Y, disc_fake_Y)\n        \n    # Calculate the gradients for generator and discriminator\n    generator_G_gradients = tape.gradient(total_gen_G_loss, \n                                              generator_G.trainable_variables)\n    generator_F_gradients = tape.gradient(total_gen_F_loss,\n                                              generator_F.trainable_variables)\n        \n    discriminator_X_gradients = tape.gradient(disc_X_loss, \n                                                  discriminator_X.trainable_variables)\n    discriminator_Y_gradients = tape.gradient(disc_Y_loss, \n                                                  discriminator_Y.trainable_variables)\n        \n    # Apply the gradients to the optimizer\n    generator_G_optimizer.apply_gradients(zip(generator_G_gradients, \n                                                  generator_G.trainable_variables))\n\n    generator_F_optimizer.apply_gradients(zip(generator_F_gradients, \n                                                  generator_F.trainable_variables))\n\n    discriminator_X_optimizer.apply_gradients(zip(discriminator_X_gradients,\n                                                      discriminator_X.trainable_variables))\n\n    discriminator_Y_optimizer.apply_gradients(zip(discriminator_Y_gradients,\n                                                      discriminator_Y.trainable_variables))","metadata":{"_uuid":"ffeddb38-abc1-4ad3-a04c-6bfcd7f7cd44","_cell_guid":"c217fa9d-d6f7-465d-9280-2935368b5edc","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-10-03T21:50:28.792252Z","iopub.status.idle":"2022-10-03T21:50:28.792988Z","shell.execute_reply.started":"2022-10-03T21:50:28.792754Z","shell.execute_reply":"2022-10-03T21:50:28.792778Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"EPOCHS = 1\n#size = train.cardinality().numpy()\n#train = train.skip(1).take(size-1)\n\nfor epoch in range(EPOCHS):\n    start = time.time()\n\n    n = 0\n    for image_X, image_Y in train:\n        train_step(image_X, image_Y)\n        if n % 10 == 0:\n            print ('.', end='')\n        n += 1\n        \n    #clear_output(wait=True)\n    # Using a consistent image (sample_dog) so that the progress of the model\n    # is clearly visible.\n    generate_images(generator_G, sample_dog)\n    generate_images(generator_F, sample_cat)\n    if (epoch + 1) % 10 == 0:\n        ckpt_save_path = ckpt_manager.save()\n        print ('Saving checkpoint for epoch {} at {}'.format(epoch+1,\n                                                             ckpt_save_path))\n        \n    print ('Time taken for epoch {} is {} sec\\n'.format(epoch + 1,\n                                                      time.time()-start))","metadata":{"_uuid":"afec2269-ad06-4ec2-8bcf-127f881bb4ec","_cell_guid":"668c26de-17e6-48e9-b02a-cb49e5e954af","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-10-03T21:50:28.794271Z","iopub.status.idle":"2022-10-03T21:50:28.795017Z","shell.execute_reply.started":"2022-10-03T21:50:28.794779Z","shell.execute_reply":"2022-10-03T21:50:28.794803Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Run the trained model on the test dataset\ninp = next(iter(train_X))\nfor k in range(BATCH_SIZE):\n    generate_images(generator_G, inp[k][tf.newaxis, ...])","metadata":{"_uuid":"6f7de337-f637-44db-a0cf-dc4ff1092068","_cell_guid":"85be2bd7-d868-4215-b95d-836536ede6da","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2022-10-03T21:50:28.796283Z","iopub.status.idle":"2022-10-03T21:50:28.797032Z","shell.execute_reply.started":"2022-10-03T21:50:28.796797Z","shell.execute_reply":"2022-10-03T21:50:28.796820Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Run the trained model on the test dataset\ninp = next(iter(train_Y))\nfor k in range(BATCH_SIZE):\n    generate_images(generator_F, inp[k][tf.newaxis, ...])","metadata":{"execution":{"iopub.status.busy":"2022-10-03T21:50:28.798341Z","iopub.status.idle":"2022-10-03T21:50:28.799103Z","shell.execute_reply.started":"2022-10-03T21:50:28.798862Z","shell.execute_reply":"2022-10-03T21:50:28.798886Z"},"trusted":true},"execution_count":null,"outputs":[]}]}