{"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":[],"dockerImageVersionId":30805,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import tensorflow as tf\nfrom tensorflow.keras import layers, models\nimport numpy as np\nimport matplotlib.pyplot as plt\n\n# Load CIFAR-10 dataset\n(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()\n\n# Normalize images to [0, 1]\nx_train = x_train.astype(\"float32\") / 255.0\nx_test = x_test.astype(\"float32\") / 255.0\nn\n# One-hot encode labels for fine-tuning\ny_train_onehot = tf.keras.utils.to_categorical(y_train, 10)\ny_test_onehot = tf.keras.utils.to_categorical(y_test, 10)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2024-12-08T08:26:58.098495Z","iopub.execute_input":"2024-12-08T08:26:58.099343Z","iopub.status.idle":"2024-12-08T08:26:59.203179Z","shell.execute_reply.started":"2024-12-08T08:26:58.099310Z","shell.execute_reply":"2024-12-08T08:26:59.202438Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"data_augmentation = tf.keras.Sequential([\n    tf.keras.layers.RandomFlip(\"horizontal\"),\n    tf.keras.layers.RandomRotation(0.2),\n    tf.keras.layers.RandomZoom(0.2)\n])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-08T08:26:59.204479Z","iopub.execute_input":"2024-12-08T08:26:59.204777Z","iopub.status.idle":"2024-12-08T08:26:59.217565Z","shell.execute_reply.started":"2024-12-08T08:26:59.204750Z","shell.execute_reply":"2024-12-08T08:26:59.216762Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Base encoder (e.g., ResNet-like small model)\ndef create_encoder():\n    inputs = layers.Input(shape=(32, 32, 3))\n    x = layers.Conv2D(32, (3, 3), activation=\"relu\", padding=\"same\")(inputs)\n    x = layers.MaxPooling2D((2, 2))(x)\n    x = layers.Conv2D(64, (3, 3), activation=\"relu\", padding=\"same\")(x)\n    x = layers.MaxPooling2D((2, 2))(x)\n    x = layers.Flatten()(x)\n    x = layers.Dense(128, activation=\"relu\")(x)  # Feature embeddings\n    return models.Model(inputs, x)\n\ndef create_projection_head(base_model):\n    inputs = layers.Input(shape=(128,))\n    x = layers.Dense(64, activation=\"relu\")(inputs)\n    x = layers.Dense(64)(x)  # Embedding space for contrastive loss\n    return models.Model(inputs, x)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-08T08:31:27.075780Z","iopub.execute_input":"2024-12-08T08:31:27.076469Z","iopub.status.idle":"2024-12-08T08:31:27.082584Z","shell.execute_reply.started":"2024-12-08T08:31:27.076435Z","shell.execute_reply":"2024-12-08T08:31:27.081801Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def contrastive_loss(projections, temperature=0.5):\n    projections = tf.nn.l2_normalize(projections, axis=1)  # Normalize projections\n    similarity_matrix = tf.matmul(projections, projections, transpose_b=True)  # Cosine similarity\n    \n    batch_size = tf.shape(projections)[0] // 2  # Half the batch size (since we have two augmentations)\n    \n    # Create labels for positive pairs (diagonal of the similarity matrix)\n    labels = tf.one_hot(tf.range(batch_size), batch_size * 2)  # Labels for positive pairs\n    labels = tf.concat([labels, labels], axis=0)  # Duplicate for both augmented views\n    \n    # logits: similarity matrix scaled by temperature\n    logits = similarity_matrix / temperature\n    \n    # Cross-entropy loss\n    loss = tf.nn.softmax_cross_entropy_with_logits(labels=labels, logits=logits)\n    \n    return tf.reduce_mean(loss)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-08T08:32:35.054608Z","iopub.execute_input":"2024-12-08T08:32:35.055483Z","iopub.status.idle":"2024-12-08T08:32:35.060723Z","shell.execute_reply.started":"2024-12-08T08:32:35.055447Z","shell.execute_reply":"2024-12-08T08:32:35.059905Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Instantiate models\nencoder = create_encoder()\nprojection_head = create_projection_head(encoder)\n\n# Adam optimizer\noptimizer = tf.keras.optimizers.Adam()\n\n# Training loop\nbatch_size = 256\nepochs = 10\n\nfor epoch in range(epochs):\n    print(f\"Epoch {epoch + 1}/{epochs}\")\n    for i in range(0, len(x_train), batch_size):\n        # Sample a batch of images\n        batch_images = x_train[i:i + batch_size]\n\n        # Apply augmentations to create positive pairs\n        augmented_1 = data_augmentation(batch_images)\n        augmented_2 = data_augmentation(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-08T08:32:35.407372Z","iopub.execute_input":"2024-12-08T08:32:35.407743Z","iopub.status.idle":"2024-12-08T08:36:05.894259Z","shell.execute_reply.started":"2024-12-08T08:32:35.407687Z","shell.execute_reply":"2024-12-08T08:36:05.893273Z"}},"outputs":[],"execution_count":null}]}