VIT_transformer.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "juHaKn0M73zj"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sFbP68C58NEm"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torchvision.transforms as transforms\n",
        "from torch.utils.data import DataLoader\n",
        "from torchvision.datasets import CIFAR10\n",
        "\n",
        "# Define the Vision Transformer (VIT) architecture\n",
        "class VIT_3c(nn.Module):\n",
        "    def __init__(self, image_width, image_height, patch_size, num_classes, dim, depth, heads, mlp_dim, dropout):\n",
        "        super(VIT_3c, self).__init__()\n",
        "\n",
        "        self.patch_embed = nn.Conv2d(3, dim, kernel_size=patch_size, stride=patch_size)\n",
        "\n",
        "        num_patches_width = image_width // patch_size\n",
        "        num_patches_height = image_height // patch_size\n",
        "\n",
        "        # num_patches = (image_size // patch_size) ** 2\n",
        "        num_patches = num_patches_width * num_patches_height\n",
        "\n",
        "        patch_dim = dim * (patch_size ** 2)\n",
        "\n",
        "        self.cls_token = nn.Parameter(torch.zeros(1, 1, dim))\n",
        "        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))\n",
        "\n",
        "        self.transformer_encoder = nn.TransformerEncoder(\n",
        "            nn.TransformerEncoderLayer(d_model=dim, nhead=heads, dim_feedforward=mlp_dim, dropout=dropout),\n",
        "            num_layers=depth\n",
        "        )\n",
        "\n",
        "        self.fc = nn.Linear(dim, num_classes)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.patch_embed(x)  # Convert image to patches\n",
        "        x = x.flatten(2).transpose(1, 2)  # Flatten patches and swap dimensions for transformer\n",
        "        cls_token = self.cls_token.expand(x.shape[0], -1, -1)  # Broadcast the class token\n",
        "        x = torch.cat((cls_token, x), dim=1)  # Concatenate the class token to patches\n",
        "        x += self.pos_embedding  # Add positional embeddings\n",
        "\n",
        "        x = self.transformer_encoder(x)  # Transformer Encoder\n",
        "\n",
        "        # Take only the class token's representation and pass through final linear layer\n",
        "        x = x[:, 0]\n",
        "        x = self.fc(x)\n",
        "        return x\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torchvision.transforms as transforms\n",
        "from torch.utils.data import DataLoader\n",
        "from torchvision.datasets import CIFAR10\n",
        "\n",
        "# Define the Vision Transformer (VIT) architecture\n",
        "class VIT_2c(nn.Module):\n",
        "    def __init__(self, image_width, image_height, patch_size, num_classes, dim, depth, heads, mlp_dim, dropout):\n",
        "        super(VIT_2c, self).__init__()\n",
        "\n",
        "        self.patch_embed = nn.Conv2d(1, dim, kernel_size=patch_size, stride=patch_size)\n",
        "\n",
        "        num_patches_width = image_width // patch_size\n",
        "        num_patches_height = image_height // patch_size\n",
        "\n",
        "        # num_patches = (image_size // patch_size) ** 2\n",
        "        num_patches = num_patches_width * num_patches_height\n",
        "\n",
        "        patch_dim = dim * (patch_size ** 2)\n",
        "\n",
        "        self.cls_token = nn.Parameter(torch.zeros(1, 1, dim))\n",
        "        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))\n",
        "\n",
        "        self.transformer_encoder = nn.TransformerEncoder(\n",
        "            nn.TransformerEncoderLayer(d_model=dim, nhead=heads, dim_feedforward=mlp_dim, dropout=dropout),\n",
        "            num_layers=depth\n",
        "        )\n",
        "\n",
        "        self.fc = nn.Linear(dim, num_classes)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.patch_embed(x)  # Convert image to patches\n",
        "        x = x.flatten(2).transpose(1, 2)  # Flatten patches and swap dimensions for transformer\n",
        "        cls_token = self.cls_token.expand(x.shape[0], -1, -1)  # Broadcast the class token\n",
        "        x = torch.cat((cls_token, x), dim=1)  # Concatenate the class token to patches\n",
        "        x += self.pos_embedding  # Add positional embeddings\n",
        "\n",
        "        x = self.transformer_encoder(x)  # Transformer Encoder\n",
        "\n",
        "        # Take only the class token's representation and pass through final linear layer\n",
        "        x = x[:, 0]\n",
        "        x = self.fc(x)\n",
        "        return x\n"
      ],
      "metadata": {
        "id": "ljtMzZBUm6Gn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# MNIST\n",
        "import torchvision.transforms as transforms\n",
        "from torchvision.datasets import MNIST\n",
        "\n",
        "# Define data transformations\n",
        "transform = transforms.Compose([\n",
        "    transforms.Resize((32, 32)), # Resize images to match the expected input size of the VIT model\n",
        "    transforms.ToTensor(), # Convert images to PyTorch tensors\n",
        "    transforms.Normalize((0.1307,), (0.3081,)) # Normalize pixel values using the mean and standard deviation of the MNIST dataset\n",
        "])\n",
        "\n",
        "# Load MNIST training dataset\n",
        "train_dataset = MNIST(root='./data', train=True, download=True, transform=transform)\n",
        "\n",
        "# Load MNIST test dataset\n",
        "test_dataset = MNIST(root='./data', train=False, download=True, transform=transform)\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UWNdcVI6l5DX",
        "outputId": "9b680ab4-8010-453b-8576-c2a477cb10b0"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz\n",
            "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to ./data/MNIST/raw/train-images-idx3-ubyte.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 9912422/9912422 [00:00<00:00, 218481461.55it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting ./data/MNIST/raw/train-images-idx3-ubyte.gz to ./data/MNIST/raw\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz\n",
            "Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz to ./data/MNIST/raw/train-labels-idx1-ubyte.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 28881/28881 [00:00<00:00, 27789789.82it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting ./data/MNIST/raw/train-labels-idx1-ubyte.gz to ./data/MNIST/raw\n",
            "\n",
            "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz\n",
            "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz to ./data/MNIST/raw/t10k-images-idx3-ubyte.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n",
            "100%|██████████| 1648877/1648877 [00:00<00:00, 66426780.48it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting ./data/MNIST/raw/t10k-images-idx3-ubyte.gz to ./data/MNIST/raw\n",
            "\n",
            "Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz\n",
            "Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to ./data/MNIST/raw/t10k-labels-idx1-ubyte.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 4542/4542 [00:00<00:00, 13165534.74it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting ./data/MNIST/raw/t10k-labels-idx1-ubyte.gz to ./data/MNIST/raw\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# HYPERPARAMETERS\n",
        "from torch.utils.data import DataLoader\n",
        "\n",
        "batch_size = 64\n",
        "\n",
        "# Create data loaders for training and test datasets\n",
        "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
        "test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "\n",
        "# Create VIT model\n",
        "model = VIT_2c(\n",
        "    image_width=32,\n",
        "    image_height=32,\n",
        "    patch_size=4,\n",
        "    num_classes=10,\n",
        "    dim=64,\n",
        "    depth=6,\n",
        "    heads=8,\n",
        "    mlp_dim=128,\n",
        "    dropout=0.1\n",
        ")\n",
        "\n",
        "# Define loss function\n",
        "criterion = nn.CrossEntropyLoss()\n",
        "\n",
        "# Define optimizer\n",
        "optimizer = optim.Adam(model.parameters(), lr=0.001)\n"
      ],
      "metadata": {
        "id": "ZWTdnHHul-de"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# train MINIST\n",
        "# Set model to training mode\n",
        "from tqdm import tqdm\n",
        "\n",
        "# train MNIST\n",
        "# Set model to training mode\n",
        "model.train()\n",
        "\n",
        "# Create a list to store the losses over time\n",
        "train_losses = []\n",
        "\n",
        "# Wrap the data loader in a tqdm object\n",
        "train_loader = tqdm(train_loader)\n",
        "\n",
        "for batch_idx, (data, target) in enumerate(train_loader):\n",
        "    # Zero out gradients from previous iteration\n",
        "    optimizer.zero_grad()\n",
        "\n",
        "    # Make predictions using VIT model\n",
        "    output = model(data)\n",
        "\n",
        "    # Calculate loss\n",
        "    loss = criterion(output, target)\n",
        "\n",
        "    # Backpropagate loss to compute gradients\n",
        "    loss.backward()\n",
        "\n",
        "    # Update model weights\n",
        "    optimizer.step()\n",
        "\n",
        "    # Append the loss to the list of losses\n",
        "    train_losses.append(loss.item())\n",
        "\n",
        "    # Set the description of the progress bar to show the loss\n",
        "    train_loader.set_description(f'Loss: {loss.item():.4f}')\n",
        "\n",
        "\n",
        "# Set model to evaluation mode\n",
        "model.eval()\n",
        "\n",
        "correct = 0\n",
        "\n",
        "with torch.no_grad():\n",
        "    for data, target in test_loader:\n",
        "        # Make predictions using VIT model\n",
        "        output = model(data)\n",
        "\n",
        "        # Get index of max log-probability as predicted class\n",
        "        pred = output.argmax(dim=1)\n",
        "\n",
        "        # Count number of correct predictions\n",
        "        correct += pred.eq(target).sum().item()\n",
        "\n",
        "# Calculate accuracy as percentage of correct predictions\n",
        "accuracy = 100. * correct / len(test_loader.dataset)\n",
        "\n",
        "print(f'Test accuracy: {accuracy:.2f}%')\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "r1Xjl3UAmoPr",
        "outputId": "c6e14914-187b-46a1-f39e-1e576e39f163"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Loss: 2.2810: 100%|██████████| 938/938 [07:25<00:00,  2.11it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Test accuracy: 9.58%\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Hyperparameters\n",
        "image_size = 32\n",
        "patch_size = 8\n",
        "num_classes = 10  # CIFAR-10 has 10 classes\n",
        "dim = 128\n",
        "depth = 6\n",
        "heads = 8\n",
        "mlp_dim = 256\n",
        "dropout = 0.1\n",
        "batch_size = 64\n",
        "epochs = 100\n",
        "learning_rate = 0.005"
      ],
      "metadata": {
        "id": "f_k5_VdPMpeA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torchvision.transforms as transforms\n",
        "from torch.utils.data import DataLoader\n",
        "from torchvision.datasets import CIFAR10\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "\n",
        "# Data preparation\n",
        "transform = transforms.Compose([\n",
        "    transforms.ToTensor(),\n",
        "    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # Normalize to [-1, 1]\n",
        "])\n",
        "\n",
        "# Load the entire CIFAR-10 dataset\n",
        "full_dataset = CIFAR10(root='./data', train=True, download=True, transform=transform)\n",
        "\n",
        "# Split the dataset into training, validation, and test sets\n",
        "num_samples = len(full_dataset)\n",
        "train_size = int(0.8 * num_samples)\n",
        "val_size = int(0.1 * num_samples)\n",
        "test_size = num_samples - train_size - val_size\n",
        "train_dataset, val_dataset, test_dataset = torch.utils.data.random_split(\n",
        "    full_dataset, [train_size, val_size, test_size])\n",
        "\n",
        "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
        "val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)\n",
        "test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)\n",
        "\n",
        "# Initialize the model, optimizer, and loss function\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "model = VIT(image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, dropout).to(device)\n",
        "optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n",
        "criterion = nn.CrossEntropyLoss()\n",
        "\n",
        "# Training loop with logging of losses\n",
        "train_losses = []\n",
        "val_losses = []\n",
        "for epoch in range(epochs):\n",
        "    model.train()\n",
        "    train_loss = 0.0\n",
        "    for images, labels in train_loader:\n",
        "        images, labels = images.to(device), labels.to(device)\n",
        "        optimizer.zero_grad()\n",
        "        outputs = model(images)\n",
        "        loss = criterion(outputs, labels)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        train_loss += loss.item() * images.size(0)\n",
        "    train_loss /= len(train_loader.dataset)\n",
        "    train_losses.append(train_loss)\n",
        "\n",
        "    # Calculate validation loss\n",
        "    model.eval()\n",
        "    val_loss = 0.0\n",
        "    with torch.no_grad():\n",
        "        for images, labels in val_loader:\n",
        "            images, labels = images.to(device), labels.to(device)\n",
        "            outputs = model(images)\n",
        "            loss = criterion(outputs, labels)\n",
        "            val_loss += loss.item() * images.size(0)\n",
        "    val_loss /= len(val_loader.dataset)\n",
        "    val_losses.append(val_loss)\n",
        "\n",
        "    print(f\"Epoch [{epoch + 1}/{epochs}], Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}\")\n",
        "\n",
        "print(\"Training finished!\")\n",
        "\n",
        "# Plotting the losses over time\n",
        "plt.plot(range(1, epochs + 1), train_losses, label='Training Loss')\n",
        "plt.plot(range(1, epochs + 1), val_losses, label='Validation Loss')\n",
        "plt.xlabel('Epochs')\n",
        "plt.ylabel('Loss')\n",
        "plt.legend()\n",
        "plt.title('Loss Over Time')\n",
        "plt.show()\n",
        "\n",
        "# Finally, evaluate the model on the test set\n",
        "model.eval()\n",
        "test_loss = 0.0\n",
        "correct = 0\n",
        "total = 0\n",
        "with torch.no_grad():\n",
        "    for images, labels in test_loader:\n",
        "        images, labels = images.to(device), labels.to(device)\n",
        "        outputs = model(images)\n",
        "        loss = criterion(outputs, labels)\n",
        "        test_loss += loss.item() * images.size(0)\n",
        "        _, predicted = torch.max(outputs.data, 1)\n",
        "        total += labels.size(0)\n",
        "        correct += (predicted == labels).sum().item()\n",
        "\n",
        "test_loss /= len(test_loader.dataset)\n",
        "test_accuracy = 100 * correct / total\n",
        "print(f\"Test Loss: {test_loss:.4f}, Test Accuracy: {test_accuracy:.2f}%\")\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rgKaMQXXMXiR",
        "outputId": "0f77b3c4-6bc3-48c0-c692-f495cd2596ed"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Files already downloaded and verified\n",
            "Epoch [1/100], Train Loss: 2.3260, Val Loss: 2.3067\n",
            "Epoch [2/100], Train Loss: 2.3056, Val Loss: 2.3049\n",
            "Epoch [3/100], Train Loss: 2.3047, Val Loss: 2.3034\n",
            "Epoch [4/100], Train Loss: 2.3047, Val Loss: 2.3037\n",
            "Epoch [5/100], Train Loss: 2.3043, Val Loss: 2.3026\n",
            "Epoch [6/100], Train Loss: 2.3044, Val Loss: 2.3031\n",
            "Epoch [7/100], Train Loss: 2.3046, Val Loss: 2.3036\n",
            "Epoch [8/100], Train Loss: 2.3041, Val Loss: 2.3050\n",
            "Epoch [9/100], Train Loss: 2.3039, Val Loss: 2.3046\n",
            "Epoch [10/100], Train Loss: 2.3039, Val Loss: 2.3032\n",
            "Epoch [11/100], Train Loss: 2.3040, Val Loss: 2.3042\n",
            "Epoch [12/100], Train Loss: 2.3040, Val Loss: 2.3040\n",
            "Epoch [13/100], Train Loss: 2.3040, Val Loss: 2.3038\n",
            "Epoch [14/100], Train Loss: 2.3040, Val Loss: 2.3071\n",
            "Epoch [15/100], Train Loss: 2.3040, Val Loss: 2.3052\n",
            "Epoch [16/100], Train Loss: 2.3041, Val Loss: 2.3028\n",
            "Epoch [17/100], Train Loss: 2.3040, Val Loss: 2.3033\n",
            "Epoch [18/100], Train Loss: 2.3040, Val Loss: 2.3049\n",
            "Epoch [19/100], Train Loss: 2.3039, Val Loss: 2.3039\n",
            "Epoch [20/100], Train Loss: 2.3042, Val Loss: 2.3075\n",
            "Epoch [21/100], Train Loss: 2.3040, Val Loss: 2.3030\n",
            "Epoch [22/100], Train Loss: 2.3038, Val Loss: 2.3033\n",
            "Epoch [23/100], Train Loss: 2.3039, Val Loss: 2.3064\n",
            "Epoch [24/100], Train Loss: 2.3042, Val Loss: 2.3052\n",
            "Epoch [25/100], Train Loss: 2.3040, Val Loss: 2.3044\n",
            "Epoch [26/100], Train Loss: 2.3042, Val Loss: 2.3036\n",
            "Epoch [27/100], Train Loss: 2.3039, Val Loss: 2.3037\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "\n",
        "# Create a 4D tensor with shape (2, 3, 4, 5)\n",
        "inputs = torch.randn(2, 3, 4, 5)\n",
        "print(f\"Original shape: {inputs.shape}\")\n",
        "\n",
        "# Reshape the tensor using *inputs.shape[2:]\n",
        "reshaped_inputs = inputs.view(-1, *inputs.shape[4:])\n",
        "print(f\"Reshaped shape: {reshaped_inputs.shape}\")\n",
        "[*inputs.shape[2:]]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2MwQTEI12GP7",
        "outputId": "6c90bd18-e825-4a12-bbef-d96d5e4eb00a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Original shape: torch.Size([2, 3, 4, 5])\n",
            "Reshaped shape: torch.Size([120])\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[4, 5]"
            ]
          },
          "metadata": {},
          "execution_count": 23
        }
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}