building_clip.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6xWqGfHyOKkp",
        "outputId": "82b03d09-1713-4d78-e077-3319585480df"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Files already downloaded and verified\n"
          ]
        }
      ],
      "source": [
        "# Check if necessary packages are installed, if not, install them\n",
        "try:\n",
        "    import torch\n",
        "    import torchvision\n",
        "    from torchvision import transforms\n",
        "    from transformers import ViTModel\n",
        "    import torch.nn as nn\n",
        "    import torch.nn.functional as F\n",
        "    import requests\n",
        "    from PIL import Image\n",
        "    from io import BytesIO\n",
        "except ImportError:\n",
        "    !pip install torch torchvision transformers\n",
        "\n",
        "    # After installing, import again\n",
        "    import torch\n",
        "    import torchvision\n",
        "    from torchvision import transforms\n",
        "    from transformers import ViTModel\n",
        "    import torch.nn as nn\n",
        "    import torch.nn.functional as F\n",
        "    import requests\n",
        "    from PIL import Image\n",
        "    from io import BytesIO\n",
        "\n",
        "# Define transformations for the dataset\n",
        "transform = transforms.Compose([\n",
        "    transforms.ToTensor(),  # Convert PIL Image to tensor\n",
        "    transforms.Resize((224, 224)),  # Resize to the input size expected by CLIP\n",
        "])\n",
        "\n",
        "# Load CIFAR-10 dataset\n",
        "train_dataset = torchvision.datasets.CIFAR10(root='./data', train=True,\n",
        "                                             download=True, transform=transform)\n",
        "\n",
        "# Custom DataLoader (optional)\n",
        "train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32,\n",
        "                                           shuffle=True, num_workers=2)\n",
        "\n",
        "# CLIP Model class\n",
        "class CLIPModel(nn.Module):\n",
        "    def __init__(self, embed_dim=512):\n",
        "        super(CLIPModel, self).__init__()\n",
        "\n",
        "        # Image backbone (ViT)\n",
        "        self.image_backbone = ViTModel.from_pretrained('google/vit-base-patch16-224-in21k')\n",
        "\n",
        "        # Text backbone (ViT, shared with image backbone)\n",
        "        self.text_backbone = self.image_backbone\n",
        "\n",
        "        # Projection MLPs\n",
        "        self.image_projection = nn.Sequential(\n",
        "            nn.Linear(768, embed_dim),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(embed_dim, embed_dim)\n",
        "        )\n",
        "\n",
        "        self.text_projection = nn.Sequential(\n",
        "            nn.Linear(768, embed_dim),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(embed_dim, embed_dim)\n",
        "        )\n",
        "\n",
        "    def forward(self, images, texts):\n",
        "        # Encode images and texts\n",
        "        image_features = self.image_backbone(images).last_hidden_state[:, 0]\n",
        "        text_features = self.text_backbone(texts).last_hidden_state[:, 0]\n",
        "\n",
        "        # Project image and text features\n",
        "        image_features = self.image_projection(image_features)\n",
        "        text_features = self.text_projection(text_features)\n",
        "\n",
        "        return image_features, text_features\n",
        "\n",
        "# Function to load an image from URL\n",
        "def load_image_from_url(url):\n",
        "    response = requests.get(url)\n",
        "    image = Image.open(BytesIO(response.content))\n",
        "    return image\n",
        "\n",
        "# Function to calculate contrastive loss\n",
        "def contrastive_loss(image_features, text_features, temperature=0.07):\n",
        "    # Normalize features\n",
        "    image_features = F.normalize(image_features, dim=-1)\n",
        "    text_features = F.normalize(text_features, dim=-1)\n",
        "\n",
        "    # Compute cosine similarity\n",
        "    logits = torch.matmul(image_features, text_features.T) / temperature\n",
        "\n",
        "    # Construct labels (0 = negative pairs, 1 = positive pairs)\n",
        "    labels = torch.arange(logits.shape[0]).to(logits.device)\n",
        "\n",
        "    # Compute cross-entropy loss\n",
        "    loss = F.cross_entropy(logits, labels)\n",
        "\n",
        "    return loss\n",
        "\n",
        "# Initialize CLIP model\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "model = CLIPModel().to(device)\n",
        "\n",
        "# Optimizer and learning rate scheduler\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)\n",
        "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)\n",
        "\n",
        "# Training loop\n",
        "def train_model(model, train_loader, optimizer, scheduler, num_epochs=10):\n",
        "    model.train()\n",
        "\n",
        "    for epoch in range(num_epochs):\n",
        "        running_loss = 0.0\n",
        "        for i, (images, _) in enumerate(train_loader):\n",
        "            images = images.to(device)\n",
        "            texts = torch.randn(images.size(0), 77).to(device)  # Random text input for demonstration\n",
        "\n",
        "            optimizer.zero_grad()\n",
        "\n",
        "            # Forward pass\n",
        "            image_features, text_features = model(images, texts)\n",
        "\n",
        "            # Calculate contrastive loss\n",
        "            loss = contrastive_loss(image_features, text_features)\n",
        "\n",
        "            # Backward pass and optimize\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "\n",
        "            running_loss += loss.item()\n",
        "\n",
        "            if (i+1) % 100 == 0:\n",
        "                print(f\"Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(train_loader)}], Loss: {running_loss / 100:.4f}\")\n",
        "                running_loss = 0.0\n",
        "\n",
        "        # Step the scheduler at the end of each epoch\n",
        "        scheduler.step()\n",
        "        print(f\"Finished Epoch [{epoch+1}/{num_epochs}], LR: {scheduler.get_last_lr()}\")\n",
        "\n",
        "# Example usage of the training loop\n",
        "if __name__ == \"__main__\":\n",
        "    train_model(model, train_loader, optimizer, scheduler, num_epochs=5)\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "AZSWgKNIOLkh"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}