{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":3649,"databundleVersionId":46718,"sourceType":"competition"}],"dockerImageVersionId":30805,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"Loading Dataset (Cr: Enrico Deccen Aristan)","metadata":{}},{"cell_type":"code","source":"import numpy as np, matplotlib.pyplot as plt, tensorflow as tf\nfrom tensorflow.keras import layers as lyr, models as mods\n\n# Load CIFAR-10 dataset\n(xtr, ytr), (xts, yts) = tf.keras.datasets.cifar10.load_data()\n\n# Input Normalization\nxtr = xtr.astype(\"float32\") / 255.0\nxts = xts.astype(\"float32\") / 255.0\n\n# Label 1-Hot Encoding\nytr_1h = tf.keras.utils.to_categorical(ytr, 10)\nyts_1h = tf.keras.utils.to_categorical(yts, 10)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-08T10:15:07.253332Z","iopub.execute_input":"2024-12-08T10:15:07.253739Z","iopub.status.idle":"2024-12-08T10:15:24.659214Z","shell.execute_reply.started":"2024-12-08T10:15:07.253689Z","shell.execute_reply":"2024-12-08T10:15:24.658138Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Data Augmenter (Cr: Enrico Deccen Aristan)","metadata":{}},{"cell_type":"code","source":"data_aug = tf.keras.Sequential([\n    lyr.RandomFlip(\"horizontal\"),\n    lyr.RandomRotation(0.2),\n    lyr.RandomZoom(0.2)\n])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-08T10:16:03.240683Z","iopub.execute_input":"2024-12-08T10:16:03.241006Z","iopub.status.idle":"2024-12-08T10:16:04.043724Z","shell.execute_reply.started":"2024-12-08T10:16:03.240976Z","shell.execute_reply":"2024-12-08T10:16:04.042549Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"SimCLR Components (Cr: Enrico Deccen Aristan)","metadata":{}},{"cell_type":"code","source":"# ResNet-like base encoder\ndef create_encoder():\n    ins = lyr.Input(shape=(32, 32, 3))\n    x = lyr.Conv2D(32, (3, 3), activation=\"relu\", padding=\"same\")(ins)\n    x = lyr.MaxPooling2D((2, 2))(x)\n    x = lyr.Conv2D(64, (3, 3), activation=\"relu\", padding=\"same\")(x)\n    x = lyr.MaxPooling2D((2, 2))(x)\n    x = lyr.Flatten()(x)\n    x = lyr.Dense(128, activation=\"relu\")(x)  # Feature embeddings\n    return mods.Model(ins, x)\n\ndef create_projection_head(base_mod):\n    ins = lyr.Input(shape=(128,))\n    x = lyr.Dense(64, activation=\"relu\")(ins)\n    x = lyr.Dense(64)(x)  # Embedding space for contrastive loss\n    return mods.Model(ins, x)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-08T10:16:07.671055Z","iopub.execute_input":"2024-12-08T10:16:07.671420Z","iopub.status.idle":"2024-12-08T10:16:07.678196Z","shell.execute_reply.started":"2024-12-08T10:16:07.671388Z","shell.execute_reply":"2024-12-08T10:16:07.677260Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Contrastive Loss Function (Cr: Enrico Deccen Aristan)","metadata":{}},{"cell_type":"code","source":"def contrastive_loss(projs, temp=0.5):\n    projs = tf.nn.l2_normalize(projs, axis=1)  # Normalize projections\n    sim_mtx = tf.matmul(projs, projs, transpose_b=True)  # Cosine similarity\n    \n    batch_size = tf.shape(projs)[0] // 2  # Halve the batch size given 2 augmentations\n    \n    # Create labels for positive pairs (diagonal of the similarity matrix)\n    lbls = tf.one_hot(tf.range(batch_size), batch_size * 2)  # Labels for positive pairs\n    lbls = tf.concat([lbls, lbls], axis=0)  # Duplicate for both augmented views\n    \n    # logits: similarity matrix scaled by temperature\n    lgts = sim_mtx / temp\n\n    loss = tf.nn.softmax_cross_entropy_with_logits(lbls, lgts)\n    \n    return tf.reduce_mean(loss)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-08T10:16:10.989721Z","iopub.execute_input":"2024-12-08T10:16:10.990477Z","iopub.status.idle":"2024-12-08T10:16:10.995409Z","shell.execute_reply.started":"2024-12-08T10:16:10.990445Z","shell.execute_reply":"2024-12-08T10:16:10.994572Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"SimCLR Training Loop (Cr: Enrico Deccen Aristan)","metadata":{}},{"cell_type":"code","source":"# Instantiate SimCLR Models\nencoder = create_encoder()\nprojection_head = create_projection_head(encoder)\n\n# Hyperparameters\noptimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)\nbatch_size = 256\nepochs = 10\n\n# Training loop\nfor epoch in range(epochs):\n    print(f\"Epoch {epoch + 1}/{epochs}\")\n    for i in range(0, len(xtr), batch_size):\n        batch_images = xtr[i:i + batch_size]\n        augmented_1 = data_aug(batch_images)\n        augmented_2 = data_aug(batch_images)\n\n        # Compute projections\n        with tf.GradientTape() as tape:\n            proj_1 = projection_head(encoder(augmented_1), training=True)\n            proj_2 = projection_head(encoder(augmented_2), training=True)\n\n            # Concatenate for NT-Xent Loss\n            projections = tf.concat([proj_1, proj_2], axis=0)\n            loss = contrastive_loss(projections)\n\n        # Apply gradients\n        gradients = tape.gradient(loss, encoder.trainable_variables + projection_head.trainable_variables)\n        optimizer.apply_gradients(zip(gradients, encoder.trainable_variables + projection_head.trainable_variables))\n\n    print(f\"Loss: {loss.numpy()}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-08T10:16:13.727819Z","iopub.execute_input":"2024-12-08T10:16:13.728442Z","iopub.status.idle":"2024-12-08T10:19:46.172429Z","shell.execute_reply.started":"2024-12-08T10:16:13.728409Z","shell.execute_reply":"2024-12-08T10:19:46.171517Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"MoCo Components","metadata":{}},{"cell_type":"code","source":"# ResNet-like base encoder\ndef moco_create_encoder():\n    ins = lyr.Input(shape=(32, 32, 3))\n    x = lyr.Conv2D(32, (3, 3), activation=\"relu\", padding=\"same\")(ins)\n    x = lyr.MaxPooling2D((2, 2))(x)\n    x = lyr.Conv2D(64, (3, 3), activation=\"relu\", padding=\"same\")(x)\n    x = lyr.MaxPooling2D((2, 2))(x)\n    x = lyr.Flatten()(x)\n    x = lyr.Dense(128, activation=\"relu\")(x)  # Feature embeddings\n    return mods.Model(ins, x)\n\ndef moco_create_projection_head(base_mod):\n    ins = lyr.Input(shape=(128,))\n    x = lyr.Dense(64, activation=\"relu\")(ins)\n    x = lyr.Dense(64)(x)  # Embedding space for contrastive loss\n    return mods.Model(ins, x)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-08T10:26:41.299407Z","iopub.execute_input":"2024-12-08T10:26:41.300158Z","iopub.status.idle":"2024-12-08T10:26:41.306674Z","shell.execute_reply.started":"2024-12-08T10:26:41.300120Z","shell.execute_reply":"2024-12-08T10:26:41.305690Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Contrastive Loss for MoCo\ndef moco_contrastive_loss(projs, memory, temp=0.5):\n    projs = tf.nn.l2_normalize(projs, axis=1) \n    memory = tf.nn.l2_normalize(memory, axis=1)\n\n    sim_mtx = tf.matmul(projs, memory, transpose_b=True) \n    \n    batch_size = tf.shape(projs)[0] // 2 \n    lbls = tf.one_hot(tf.range(batch_size), batch_size * 2)\n    lbls = tf.concat([lbls, lbls], axis=0)\n    \n    lgts = sim_mtx / temp\n    loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(lbls, lgts))\n\n    return loss\n\n# MoCo-specific momentum encoder update function\ndef moco_update_momentum_encoder(encoder, momentum_encoder, momentum=0.99):\n    for var, momentum_var in zip(encoder.trainable_variables, momentum_encoder.trainable_variables):\n        momentum_var.assign(momentum * momentum_var + (1.0 - momentum) * var)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-08T10:26:47.198549Z","iopub.execute_input":"2024-12-08T10:26:47.198932Z","iopub.status.idle":"2024-12-08T10:26:47.206055Z","shell.execute_reply.started":"2024-12-08T10:26:47.198851Z","shell.execute_reply":"2024-12-08T10:26:47.204622Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"MoCo Training Loop","metadata":{}},{"cell_type":"code","source":"# Instantiate MoCo Models\nencoder = moco_create_encoder()\nprojection_head = moco_create_projection_head(encoder)\n\n# Initialize the momentum encoder (target encoder)\nmomentum_encoder = moco_create_encoder()\nmomentum_projection_head = moco_create_projection_head(momentum_encoder)\n\n# Hyperparameters\noptimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)\nbatch_size = 256\nepochs = 10\nmomentum = 0.99  # Momentum for the momentum encoder\n\n# Memory bank (used to store the keys from the momentum encoder)\nmemory_bank = tf.Variable(tf.zeros((batch_size, 64)), trainable=False)  # Adjust dimensions as necessary\n\n# MoCo Training Loop\nfor epoch in range(epochs):\n    print(f\"Epoch {epoch + 1}/{epochs}\")\n    for i in range(0, len(xtr), batch_size):\n        batch_images = xtr[i:i + batch_size]\n        augmented_1 = data_aug(batch_images)\n        augmented_2 = data_aug(batch_images)\n\n        # Compute projections\n        with tf.GradientTape() as tape:\n            proj_1 = projection_head(encoder(augmented_1), training=True)\n            proj_2 = projection_head(encoder(augmented_2), training=True)\n\n            # Concatenate projections from two augmentations\n            projections = tf.concat([proj_1, proj_2], axis=0)\n\n            # Get memory bank (keys from momentum encoder)\n            momentum_proj_1 = momentum_projection_head(momentum_encoder(augmented_1), training=False)\n            momentum_proj_2 = momentum_projection_head(momentum_encoder(augmented_2), training=False)\n            memory_proj = tf.concat([momentum_proj_1, momentum_proj_2], axis=0)\n            \n            # Compute MoCo contrastive loss\n            loss = moco_contrastive_loss(projections, memory_proj)\n\n        # Apply gradients\n        gradients = tape.gradient(loss, encoder.trainable_variables + projection_head.trainable_variables)\n        optimizer.apply_gradients(zip(gradients, encoder.trainable_variables + projection_head.trainable_variables))\n\n        # Update momentum encoder\n        moco_update_momentum_encoder(encoder, momentum_encoder, momentum=momentum)\n\n    print(f\"Loss: {loss.numpy()}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-08T10:26:49.599708Z","iopub.execute_input":"2024-12-08T10:26:49.600037Z","iopub.status.idle":"2024-12-08T10:30:59.909400Z","shell.execute_reply.started":"2024-12-08T10:26:49.600006Z","shell.execute_reply":"2024-12-08T10:30:59.908505Z"}},"outputs":[],"execution_count":null}]}