{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":21755,"databundleVersionId":1475600,"sourceType":"competition"},{"sourceId":12363156,"sourceType":"datasetVersion","datasetId":7794878}],"dockerImageVersionId":31040,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"* too lazy to submit predictions,just sharing knowledge ","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport matplotlib.pyplot as plt\nimport tensorflow as tf\nfrom tensorflow.keras.utils import register_keras_serializable\nimport base64\nfrom IPython.display import HTML\nimport os\nfrom datetime import datetime","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-07-03T14:24:17.534325Z","iopub.execute_input":"2025-07-03T14:24:17.535200Z","iopub.status.idle":"2025-07-03T14:24:35.274844Z","shell.execute_reply.started":"2025-07-03T14:24:17.535151Z","shell.execute_reply":"2025-07-03T14:24:35.273788Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def gif_to_html(gif_path):\n    with open(gif_path, 'rb') as f:\n        encoded = base64.b64encode(f.read()).decode('utf-8')\n    return f'<img src=\"data:image/gif;base64,{encoded}\">'\nHTML(gif_to_html(\"/kaggle/input/photo-to-monet-gif/photo to monet cut gan gif.gif\"))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-03T14:24:35.276803Z","iopub.execute_input":"2025-07-03T14:24:35.277474Z","iopub.status.idle":"2025-07-03T14:24:35.498580Z","shell.execute_reply.started":"2025-07-03T14:24:35.277445Z","shell.execute_reply":"2025-07-03T14:24:35.497469Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"path_photo = \"/kaggle/input/gan-getting-started/photo_jpg\"\npath_monet = \"/kaggle/input/gan-getting-started/monet_jpg\"\nphoto_dir = tf.data.Dataset.list_files(path_photo + \"/*\",shuffle=False)\nmonet_dir = tf.data.Dataset.list_files(path_monet + \"/*\",shuffle=False)\n#load img\ndef decode_img(path):\n    path = tf.io.read_file(path)\n    img = tf.io.decode_png(path,channels=3)\n    img = tf.image.resize(img,(256,256),method='area')\n    img = tf.cast(img,tf.float32)\n    img = img / 127.5 - 1.0\n    return img \n\ndef denorm(img):\n    return (img + 1) / 2 #[-1,1] -> [0,1]\n\nphoto = photo_dir.map(decode_img,num_parallel_calls=tf.data.AUTOTUNE)\nmonet = monet_dir.map(decode_img,num_parallel_calls=tf.data.AUTOTUNE)\ndataset = tf.data.Dataset.zip((photo,monet))\nval_dataset = dataset.take(10)\ntrain_dataset = dataset.skip(10)\n\ntrain_dataset = train_dataset.batch(16,drop_remainder=True).cache().prefetch(tf.data.AUTOTUNE)\nval_dataset = val_dataset.batch(5,drop_remainder=True).cache().prefetch(tf.data.AUTOTUNE)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-03T14:24:35.499593Z","iopub.execute_input":"2025-07-03T14:24:35.500028Z","iopub.status.idle":"2025-07-03T14:24:42.408013Z","shell.execute_reply.started":"2025-07-03T14:24:35.499994Z","shell.execute_reply":"2025-07-03T14:24:42.406973Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"size = len(list(train_dataset.take(1).unbatch()))\nrand_index = np.random.randint(0, size - 1)\nsize_val = len(list(val_dataset.take(1).unbatch()))\nrand_index_val = np.random.randint(0, size_val - 1)\n\nfor (photo_train, monet_train), (photo_val, monet_val) in zip(train_dataset.take(1), val_dataset.take(1)):\n    plt.figure(figsize=(10, 2))\n\n    plt.subplot(1, 4, 1)\n    plt.imshow(denorm(photo_train[rand_index]))\n    plt.title(\"photo Train\", fontsize=8)\n\n    plt.subplot(1, 4, 2)\n    plt.imshow(denorm(monet_train[rand_index]))\n    plt.title(\"monet Train\", fontsize=8)\n\n    plt.subplot(1, 4, 3)\n    plt.imshow(denorm(photo_val[rand_index_val]))\n    plt.title(\"photo Val\", fontsize=8)\n\n    plt.subplot(1, 4, 4)\n    plt.imshow(denorm(monet_val[rand_index_val]))\n    plt.title(\"monet Val\", fontsize=8)\n\n    plt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-03T14:25:00.396426Z","iopub.execute_input":"2025-07-03T14:25:00.396818Z","iopub.status.idle":"2025-07-03T14:25:01.529346Z","shell.execute_reply.started":"2025-07-03T14:25:00.396789Z","shell.execute_reply":"2025-07-03T14:25:01.528400Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def DownSample(inputs, filters,kernel_size=4, strides=2, activation=None,index=0 ,inst_norm=False):\n    x = tf.keras.layers.SpectralNormalization(tf.keras.layers.Conv2D(filters=filters,\n                  kernel_size=kernel_size,\n                  strides=strides,\n                  padding='same',name=f'downsample_{index}'),power_iterations=5)(inputs)\n    if inst_norm:\n        x = tf.keras.layers.GroupNormalization(groups=filters)(x)\n    if activation == 'relu':\n        x = tf.keras.layers.Activation('relu')(x)\n    if activation == 'leakyrelu':\n        x = tf.keras.layers.LeakyReLU(alpha=0.2)(x)\n\n\n    return x\n\ndef UpSample(inputs,\n             paired_inputs,\n             filters,         \n             kernel_size=4,\n             strides=2,\n             activation=None,\n             inst_norm=False):\n    # Concatenate BEFORE upsampling\n    x = tf.keras.layers.Concatenate()([inputs, paired_inputs])\n\n    # Transposed convolution layer.\n    x = tf.keras.layers.Conv2DTranspose(filters=filters,\n                                        kernel_size=kernel_size,\n                                        strides=strides,\n                                        padding='same')(x)\n    if inst_norm:\n        x = tf.keras.layers.GroupNormalization(groups=filters)(x)\n    if activation == 'relu':\n        x = tf.keras.layers.Activation('relu')(x)\n    elif activation == 'leakyrelu':\n        x = tf.keras.layers.LeakyReLU(alpha=0.2)(x)\n    return x\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-03T13:01:14.693132Z","iopub.execute_input":"2025-07-03T13:01:14.693407Z","iopub.status.idle":"2025-07-03T13:01:14.700268Z","shell.execute_reply.started":"2025-07-03T13:01:14.693386Z","shell.execute_reply":"2025-07-03T13:01:14.699682Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from tensorflow.keras import layers, Model\n\ndef ResBlock(x, filters):\n    shortcut = x\n    x = layers.Conv2D(filters, kernel_size=3, padding='same', use_bias=False)(x)\n    x = layers.BatchNormalization()(x)\n    x = layers.ReLU()(x)\n    x = layers.Conv2D(filters, kernel_size=3, padding='same', use_bias=False)(x)\n    x = layers.BatchNormalization()(x)\n    x = layers.Add()([x, shortcut])\n    return x\n\ndef build_generator(input_shape=(256, 256, 3)):\n    inputs = layers.Input(shape=input_shape)\n\n    # ---------- Encoder / Downsampling ----------\n    e1 = DownSample(inputs, 64, kernel_size=4, strides=1, activation='leakyrelu', inst_norm=False, index=1)  # No downsample, just first conv\n    e2 = DownSample(e1, 128, kernel_size=4, strides=2, activation='leakyrelu', inst_norm=True, index=2)\n    e3 = DownSample(e2, 192, kernel_size=4, strides=2, activation='leakyrelu', inst_norm=True, index=3)\n    e4 = DownSample(e3, 256, kernel_size=4, strides=2, activation='leakyrelu', inst_norm=True, index=4)\n\n    # ---------- ResNet Bottleneck ----------\n    x = e4\n    res_feats = []\n    for i in range(9):\n        x = ResBlock(x, 256)\n        if i in [2, 5, 8]:  # Choose 3 ResBlocks as deeper NCE layers\n            res_feats.append(x)\n\n    # ---------- Decoder / Upsampling ----------\n    d1 = UpSample(x, e4, 256, kernel_size=4, strides=2, activation='leakyrelu', inst_norm=True)\n    d2 = UpSample(d1, e3, 192, kernel_size=4, strides=2, activation='leakyrelu', inst_norm=True)\n    d3 = UpSample(d2, e2, 128, kernel_size=4, strides=2, activation='leakyrelu', inst_norm=True)\n    d4 = UpSample(d3, e1, 64, kernel_size=4, strides=1, activation='leakyrelu', inst_norm=True)\n\n    outputs = layers.Conv2DTranspose(3, kernel_size=4, strides=1, padding='same', activation='tanh')(d4)\n\n    # ---------- NCE Features ----------\n    # Total: 5 layers → (e2, e3, e4, res3, res6)\n    feats_for_NCE = [e2, e3, e4] + res_feats\n\n    return Model(inputs, feats_for_NCE + [outputs])\n    \ndef build_discriminator(input_shape, name=None):\n    Inputs = tf.keras.layers.Input(shape=input_shape)\n    e1 = DownSample(Inputs,64,kernel_size=4,activation='leakyrelu',strides=2,inst_norm=False)\n    e2 = DownSample(e1,128,kernel_size=4,activation='leakyrelu',strides=2,inst_norm=True)\n    e3 = DownSample(e2,256,kernel_size=4,activation='leakyrelu',strides=2,inst_norm=True)\n    e4 = DownSample(e3,512,kernel_size=4,activation='leakyrelu',strides=2,inst_norm=True)\n    PatchGAN = tf.keras.layers.Conv2D(1, kernel_size=4, strides=1, padding='same')(e4)\n    discriminator = tf.keras.models.Model(Inputs, PatchGAN, name=name)\n    return discriminator","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-03T13:01:24.772816Z","iopub.execute_input":"2025-07-03T13:01:24.773082Z","iopub.status.idle":"2025-07-03T13:01:24.784721Z","shell.execute_reply.started":"2025-07-03T13:01:24.773063Z","shell.execute_reply":"2025-07-03T13:01:24.784109Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"G_X = build_generator((256,256,3))\nD_X = build_discriminator((256,256,3))\nG_X.summary()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-03T13:01:33.606294Z","iopub.execute_input":"2025-07-03T13:01:33.606968Z","iopub.status.idle":"2025-07-03T13:01:35.791212Z","shell.execute_reply.started":"2025-07-03T13:01:33.606941Z","shell.execute_reply":"2025-07-03T13:01:35.790454Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class VisualizationCallback(tf.keras.callbacks.Callback):\n    def __init__(self, trainer, test_input_a, test_input_b, save_freq=1, save_dir=\"visualizations\", num_images=4):\n        super().__init__()\n        self.trainer = trainer\n        self.test_input_a = test_input_a\n        self.test_input_b = test_input_b\n        self.num_images = num_images\n        self.save_freq = save_freq\n\n        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')\n        self.save_dir = os.path.join(save_dir, f\"run_{timestamp}\")\n        os.makedirs(self.save_dir, exist_ok=True)\n\n    def on_epoch_end(self, epoch, logs=None):\n        if (epoch + 1) % self.save_freq != 0:\n            return\n\n        fake_images = self.trainer.generator(self.test_input_a, training=False)[-1]\n\n        fig, axes = plt.subplots(3, self.num_images, figsize=(self.num_images * 3, 9))\n\n        for i in range(self.num_images):\n            # Row 1: Input A\n            img_a = (self.test_input_a[i] + 1.0) / 2.0\n            img_a = tf.clip_by_value(img_a, 0.0, 1.0)\n            axes[0, i].imshow(img_a.numpy())\n            axes[0, i].axis('off')\n            axes[0, i].set_title(f\"Input A {i+1}\")\n\n            # Row 2: Generated fake B\n            img_b_fake = (fake_images[i] + 1.0) / 2.0\n            img_b_fake = tf.clip_by_value(img_b_fake, 0.0, 1.0)\n            axes[1, i].imshow(img_b_fake.numpy())\n            axes[1, i].axis('off')\n            axes[1, i].set_title(f\"Fake B {i+1}\")\n\n            # Row 3: Ground Truth real B\n            img_b_real = (self.test_input_b[i] + 1.0) / 2.0\n            img_b_real = tf.clip_by_value(img_b_real, 0.0, 1.0)\n            axes[2, i].imshow(img_b_real.numpy())\n            axes[2, i].axis('off')\n            axes[2, i].set_title(f\"Real B {i+1}\")\n\n        plt.tight_layout()\n        save_path = os.path.join(self.save_dir, f\"epoch_{epoch+1:04d}.png\")\n        plt.savefig(save_path)\n        plt.close()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-03T13:01:51.400723Z","iopub.execute_input":"2025-07-03T13:01:51.401263Z","iopub.status.idle":"2025-07-03T13:01:51.409668Z","shell.execute_reply.started":"2025-07-03T13:01:51.401243Z","shell.execute_reply":"2025-07-03T13:01:51.409025Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class PatchNCELoss(tf.keras.layers.Layer):\n    def __init__(self, temperature=0.07, **kwargs):\n        super().__init__(**kwargs)\n        self.temperature = temperature\n        self.cross_entropy_loss = tf.keras.losses.SparseCategoricalCrossentropy(\n            from_logits=True, reduction=tf.keras.losses.Reduction.NONE\n        )\n\n    def call(self, f_q, f_k):\n        # f_q: query features (fake), f_k: key features (real)\n        # Input shape: (B, S, C) where S is number of patches, C is feature dim\n        \n        # L2 normalize\n        f_q = tf.nn.l2_normalize(f_q, axis=-1)  # (B, S, C)\n        f_k = tf.nn.l2_normalize(f_k, axis=-1)  # (B, S, C)\n\n        batch_size = tf.shape(f_q)[0]\n        num_patches = tf.shape(f_q)[1]\n        \n        # Transpose for easier computation: (B, C, S)\n        f_q_t = tf.transpose(f_q, [0, 2, 1])  # (B, C, S)\n        f_k_t = tf.transpose(f_k, [0, 2, 1])  # (B, C, S)\n\n        # Calculate positive similarities: v * v+ (element-wise multiplication and sum)\n        # l_pos shape: (B, S, 1)\n        l_pos = tf.reduce_sum(f_k * f_q, axis=-1, keepdims=True)  # (B, S, 1)\n\n        # Calculate negative similarities: v * v- (batch matrix multiplication)\n        # f_q: (B, S, C), f_k_t: (B, C, S) -> l_neg: (B, S, S)\n        l_neg = tf.matmul(f_q, f_k_t)  # (B, S, S)\n\n        # Remove diagonal entries (self-similarities) by setting them to -inf\n        identity_matrix = tf.eye(num_patches, dtype=tf.float32)  # (S, S)\n        identity_matrix = tf.expand_dims(identity_matrix, 0)  # (1, S, S)\n        l_neg = tf.where(identity_matrix > 0, -float('inf'), l_neg)\n\n        # Concatenate positive and negative logits: (B, S, S+1)\n        logits = tf.concat([l_pos, l_neg], axis=2) / self.temperature\n\n        # Flatten to (B*S, S+1) for cross-entropy loss\n        predictions = tf.reshape(logits, [-1, tf.shape(logits)[-1]])\n        \n        # Targets are always 0 (first column is positive)\n        targets = tf.zeros(batch_size * num_patches, dtype=tf.int32)\n        \n        loss = self.cross_entropy_loss(targets, predictions)\n        return tf.reduce_mean(loss)\n\nclass PatchSampleMLP(tf.keras.layers.Layer):\n    def __init__(self, units=256, num_patches=256, **kwargs):\n        super().__init__(**kwargs)\n        self.units = units\n        self.num_patches = num_patches\n        self.mlp = tf.keras.Sequential([\n            tf.keras.layers.Dense(self.units, activation='relu'),\n            tf.keras.layers.Dense(self.units)\n        ])\n\n    def call(self, feat, training=True):\n        # Flatten spatial dimensions\n        B, H, W, C = tf.shape(feat)[0], tf.shape(feat)[1], tf.shape(feat)[2], tf.shape(feat)[3]\n        feat = tf.reshape(feat, [B, H*W, C])\n\n        num_patches = tf.minimum(self.num_patches, H * W)\n        \n        # Sample patches consistently across batch\n        if training:\n            rand_idx = tf.random.shuffle(tf.range(H * W))[:num_patches]\n        else:\n            # Use deterministic sampling for validation\n            rand_idx = tf.range(num_patches)\n            \n        sampled_feat = tf.gather(feat, rand_idx, axis=1)  # (B, num_patches, C)\n\n        proj_feat = self.mlp(sampled_feat, training=training)  # (B, num_patches, units)\n        proj_feat = tf.nn.l2_normalize(proj_feat, axis=-1)\n        return proj_feat\n\n\nclass GANTrainer(tf.keras.Model):\n    def __init__(self, generator, discriminator, **kwargs):\n        super().__init__(**kwargs)\n        self.generator = generator\n        self.discriminator = discriminator\n\n        # MLPs for each layer of features\n        self.mlp_list = [\n            PatchSampleMLP(units=256, num_patches=256),\n            PatchSampleMLP(units=256, num_patches=256),\n            PatchSampleMLP(units=256, num_patches=256),\n            PatchSampleMLP(units=256, num_patches=256),\n            PatchSampleMLP(units=256, num_patches=256),\n        ]\n\n        # PatchNCE Loss\n        self.patchnce_loss_fn = PatchNCELoss()\n\n    def compile(self, gen_optimizer, disc_optimizer, mlp_optimizer, **kwargs):\n        super().compile(**kwargs)\n        self.gen_optimizer = gen_optimizer\n        self.disc_optimizer = disc_optimizer\n        self.mlp_optimizer = mlp_optimizer\n        self.loss_metric = tf.keras.metrics.Mean(name=\"total_gen_loss\")\n        self.disc_loss_metric = tf.keras.metrics.Mean(name=\"disc_loss\")\n        self.nce_loss_metric = tf.keras.metrics.Mean(name=\"nce_loss\")\n\n    def lsgan_discriminator_loss(self, real_output, fake_output):\n        real_loss = tf.reduce_mean(tf.square(real_output - 1.0))\n        fake_loss = tf.reduce_mean(tf.square(fake_output))\n        return 0.5 * (real_loss + fake_loss)\n\n    def lsgan_generator_loss(self, fake_output):\n        return tf.reduce_mean(tf.square(fake_output - 1.0))\n\n    def calculate_NCE_loss(self, src, tgt, training=True):\n        \"\"\"\n        Calculate NCE loss following the reference implementation\n        src: source image (real_A)\n        tgt: target image (fake_B)\n        \"\"\"\n        # Extract features from both images using the generator encoder\n        # feat_k: features from source (real_A)\n        feat_k_output = self.generator(src, training=training)\n        feat_k = feat_k_output[:-1]  # Feature maps only (excluding final output)\n        \n        # feat_q: features from target (fake_B)  \n        feat_q_output = self.generator(tgt, training=training)\n        feat_q = feat_q_output[:-1]  # Feature maps only\n        \n        # Optional: flip equivariance (set to False if not needed)\n        flip_equivariance = False\n        if flip_equivariance:\n            feat_q = [tf.reverse(fq, axis=[2]) for fq in feat_q]  # Flip width dimension\n        \n        total_nce_loss = 0.0\n        n_layers = len(feat_k)\n        \n        for layer_idx, (f_k, f_q, mlp) in enumerate(zip(feat_k, feat_q, self.mlp_list)):\n            # First, sample patches from feat_k and get sample_ids\n            feat_k_pool, sample_ids = self.sample_patches_with_ids(f_k, mlp, training)\n            \n            # Use the same sample_ids for feat_q to ensure correspondence\n            feat_q_pool = self.sample_patches_with_ids(f_q, mlp, training, sample_ids)\n            \n            # Compute NCE loss for this layer\n            loss = self.patchnce_loss_fn(feat_q_pool, feat_k_pool)\n            total_nce_loss += loss\n            \n        return total_nce_loss / n_layers\n    \n    def sample_patches_with_ids(self, feat, mlp, training=True, sample_ids=None):\n        \"\"\"\n        Sample patches from feature maps, optionally using provided sample_ids\n        Returns: (projected_features, sample_ids) or just projected_features if sample_ids provided\n        \"\"\"\n        # Flatten spatial dimensions\n        B, H, W, C = tf.shape(feat)[0], tf.shape(feat)[1], tf.shape(feat)[2], tf.shape(feat)[3]\n        feat_flat = tf.reshape(feat, [B, H*W, C])\n        \n        num_patches = tf.minimum(mlp.num_patches, H * W)\n        \n        if sample_ids is None:\n            # Generate new sample_ids\n            if training:\n                sample_ids = tf.random.shuffle(tf.range(H * W))[:num_patches]\n            else:\n                sample_ids = tf.range(num_patches)\n            \n            # Sample patches using the generated ids\n            sampled_feat = tf.gather(feat_flat, sample_ids, axis=1)  # (B, num_patches, C)\n            \n            # Project features using MLP\n            proj_feat = mlp.mlp(sampled_feat, training=training)  # (B, num_patches, units)\n            proj_feat = tf.nn.l2_normalize(proj_feat, axis=-1)\n            \n            return proj_feat, sample_ids\n        else:\n            # Use provided sample_ids\n            sampled_feat = tf.gather(feat_flat, sample_ids, axis=1)  # (B, num_patches, C)\n            \n            # Project features using MLP\n            proj_feat = mlp.mlp(sampled_feat, training=training)  # (B, num_patches, units)\n            proj_feat = tf.nn.l2_normalize(proj_feat, axis=-1)\n            \n            return proj_feat\n\n    @tf.function\n    def train_step(self, data):\n        real_A, real_B = data\n    \n        with tf.GradientTape(persistent=True) as tape:\n            # Forward pass: real_A -> fake_B\n            real_A_output = self.generator(real_A, training=True)\n            fake_B = real_A_output[-1]  # Final output\n    \n            # Discriminator outputs\n            real_out = self.discriminator(real_B, training=True)\n            fake_out = self.discriminator(fake_B, training=True)\n            disc_loss = self.lsgan_discriminator_loss(real_out, fake_out)\n    \n            # Generator adversarial loss\n            gen_gan_loss = self.lsgan_generator_loss(fake_out)\n            \n            # NCE loss: encourage correspondence between real_A and fake_B features\n            # Following reference: src=real_A, tgt=fake_B\n            nce_loss = self.calculate_NCE_loss(real_A, fake_B, training=True)\n            \n            # Identity loss (optional - use for CUT, skip for FastCUT)\n            identity_loss = 0.0\n            if True:  # Set to False for FastCUT\n                real_B_output = self.generator(real_B, training=True)\n                idt_B = real_B_output[-1]\n                identity_loss = tf.reduce_mean(tf.abs(real_B - idt_B))\n            \n            # Total generator loss\n            # CUT: lambda_NCE=1, lambda_identity=0.5\n            # FastCUT: lambda_NCE=10, lambda_identity=0\n            lambda_nce = 1.0\n            lambda_identity = 0.5\n            total_gen_loss = gen_gan_loss + lambda_nce * nce_loss + lambda_identity * identity_loss\n    \n        # Update discriminator\n        disc_grads = tape.gradient(disc_loss, self.discriminator.trainable_variables)\n        self.disc_optimizer.apply_gradients(zip(disc_grads, self.discriminator.trainable_variables))\n    \n        # Update generator\n        gen_grads = tape.gradient(total_gen_loss, self.generator.trainable_variables)\n        self.gen_optimizer.apply_gradients(zip(gen_grads, self.generator.trainable_variables))\n    \n        # Update MLPs (only on NCE loss)\n        mlp_vars = []\n        for mlp in self.mlp_list:\n            mlp_vars += mlp.trainable_variables\n        mlp_grads = tape.gradient(nce_loss, mlp_vars)\n        self.mlp_optimizer.apply_gradients(zip(mlp_grads, mlp_vars))\n    \n        # Update metrics\n        self.loss_metric.update_state(total_gen_loss)\n        self.disc_loss_metric.update_state(disc_loss)\n        self.nce_loss_metric.update_state(nce_loss)\n    \n        return {\n            \"total_gen_loss\": self.loss_metric.result(),\n            \"disc_loss\": self.disc_loss_metric.result(),\n            \"nce_loss\": self.nce_loss_metric.result(),\n            \"identity_loss\": identity_loss,\n        }\n\n    @property\n    def metrics(self):\n        return [self.loss_metric, self.disc_loss_metric, self.nce_loss_metric]\n\n\n# Usage\ntrainer = GANTrainer(G_X, D_X)\n\ntrainer.compile(\n    gen_optimizer=tf.keras.optimizers.Adam(2e-4, beta_1=0.5),\n    disc_optimizer=tf.keras.optimizers.Adam(2e-4, beta_1=0.5),\n    mlp_optimizer=tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n)\n\n# Create test samples for visualization (batch of images from domain A and B)\ntest_input_a = next(iter(val_dataset))[0]  # Domain A sample\ntest_input_b = next(iter(val_dataset))[1]  # Domain B sample\n\n# Initialize visualization callback\nvis_callback = VisualizationCallback(\n    trainer=trainer,\n    test_input_a=test_input_a,\n    test_input_b=test_input_b,\n    save_freq=1  # Save images every epoch\n)\n\n# Train with visualization callback\ntrainer.fit(\n    train_dataset,\n    epochs=50,\n    callbacks=[vis_callback]\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-03T13:02:18.847485Z","iopub.execute_input":"2025-07-03T13:02:18.848044Z","iopub.status.idle":"2025-07-03T13:35:55.378848Z","shell.execute_reply.started":"2025-07-03T13:02:18.848019Z","shell.execute_reply":"2025-07-03T13:35:55.377911Z"}},"outputs":[],"execution_count":null}]}