simclr.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5b45e933be2198bd",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# SimCLR Implementation and Evaluation on CIFAR-10\n",
    "\n",
    "This notebook implements the SimCLR algorithm, trains it on the CIFAR-10 dataset, and evaluates the learned representations using Linear Probing and K-Nearest Neighbors (KNN) classification.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d3a06b9e367185d3",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-09-27T09:30:47.495463920Z",
     "start_time": "2023-09-27T09:30:46.320579552Z"
    }
   },
   "outputs": [],
   "source": [
    "# Import necessary libraries\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "from torch.utils.data import DataLoader\n",
    "from sklearn.neighbors import KNeighborsClassifier\n",
    "import numpy as np\n",
    "import torch.nn.functional as F\n",
    "import torchvision.models as models\n",
    "\n",
    "# Importing necessary libraries and modules for the implementation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b72739b",
   "metadata": {},
   "source": [
    "### Execution Timers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "54d447ac",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-27T09:30:47.498344303Z",
     "start_time": "2023-09-27T09:30:47.495919475Z"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "# Flag to enable or disable timers\n",
    "enable_timers = True\n",
    "\n",
    "import time\n",
    "\n",
    "class Timer:\n",
    "    def __enter__(self):\n",
    "        if enable_timers:\n",
    "            self.start = time.time()\n",
    "        return self\n",
    "\n",
    "    def __exit__(self, *args):\n",
    "        if enable_timers:\n",
    "            self.end = time.time()\n",
    "            self.interval = self.end - self.start\n",
    "            print(f\"Elapsed time: {self.interval:.2f} seconds\")\n",
    "    \n",
    "\n",
    "# Importing necessary libraries and modules for the implementation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b017aff08498197",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Load CIFAR-10 Dataset\n",
    "\n",
    "Load the CIFAR-10 training and test datasets.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "184d173ca4466895",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-09-27T09:30:48.004081941Z",
     "start_time": "2023-09-27T09:30:47.497687238Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "from data_aug.contrastive_learning_dataset import ContrastiveLearningDataset\n",
    "\n",
    "dataset = ContrastiveLearningDataset(root_folder='data')\n",
    "train_dataset = dataset.get_dataset('cifar10', 2)\n",
    "train_loader = torch.utils.data.DataLoader(\n",
    "        train_dataset, batch_size=512, shuffle=True,\n",
    "        num_workers=8, pin_memory=True, drop_last=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [],
   "metadata": {
    "collapsed": false
   },
   "id": "52453cbfe9cae13a"
  },
  {
   "cell_type": "markdown",
   "id": "f50c4c48b88d94ba",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Define SimCLR Encoder and Projection Head\n",
    "\n",
    "Create the encoder model and projection head using ResNet18 as the base architecture.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5ee2abd1db6cba59",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-09-27T09:30:49.841557381Z",
     "start_time": "2023-09-27T09:30:49.837102964Z"
    }
   },
   "outputs": [],
   "source": [
    "class ResNetSimCLR(nn.Module):\n",
    "\n",
    "    def __init__(self, base_model, out_dim):\n",
    "        super(ResNetSimCLR, self).__init__()\n",
    "        self.resnet_dict = {\"resnet18\": models.resnet18(pretrained=False, num_classes=out_dim),\n",
    "                            \"resnet50\": models.resnet50(pretrained=False, num_classes=out_dim)}\n",
    "\n",
    "        self.backbone = self._get_basemodel(base_model)\n",
    "        dim_mlp = self.backbone.fc.in_features\n",
    "\n",
    "        # add mlp projection head\n",
    "        self.backbone.fc = nn.Sequential(nn.Linear(dim_mlp, dim_mlp), nn.ReLU(), self.backbone.fc)\n",
    "\n",
    "    def _get_basemodel(self, model_name):\n",
    "        try:\n",
    "            model = self.resnet_dict[model_name]\n",
    "        except:\n",
    "            raise (\"Invalid model name. Check the config file and pass one of: resnet18 or resnet50\")\n",
    "\n",
    "        return model\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.backbone(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d149f22bdda4adaf",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Define Contrastive Loss\n",
    "\n",
    "Implement the contrastive loss function used by SimCLR.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "577b5e81f10c8164",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-09-27T09:30:50.988597175Z",
     "start_time": "2023-09-27T09:30:50.982323266Z"
    }
   },
   "outputs": [],
   "source": [
    "def info_nce_loss(features, temperature=0.5):\n",
    "        batch_size = features.shape[0] // 2 # 2 views per batch\n",
    "        \n",
    "        labels = torch.cat([torch.arange(batch_size) for i in range(2)], dim=0)\n",
    "        labels = (labels.unsqueeze(0) == labels.unsqueeze(1)).float()\n",
    "        device = (torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu'))\n",
    "        labels = labels.to(device)\n",
    "\n",
    "        features = F.normalize(features, dim=1)\n",
    "\n",
    "        similarity_matrix = torch.matmul(features, features.T)\n",
    "        # assert similarity_matrix.shape == (\n",
    "        #     self.args.n_views * self.args.batch_size, self.args.n_views * self.args.batch_size)\n",
    "        # assert similarity_matrix.shape == labels.shape\n",
    "\n",
    "        # discard the main diagonal from both: labels and similarities matrix\n",
    "        mask = torch.eye(labels.shape[0], dtype=torch.bool).to(device)\n",
    "        labels = labels[~mask].view(labels.shape[0], -1)\n",
    "        similarity_matrix = similarity_matrix[~mask].view(similarity_matrix.shape[0], -1)\n",
    "        # assert similarity_matrix.shape == labels.shape\n",
    "\n",
    "        # select and combine multiple positives\n",
    "        positives = similarity_matrix[labels.bool()].view(labels.shape[0], -1)\n",
    "\n",
    "        # select only the negatives\n",
    "        negatives = similarity_matrix[~labels.bool()].view(similarity_matrix.shape[0], -1)\n",
    "\n",
    "        logits = torch.cat([positives, negatives], dim=1)\n",
    "        labels = torch.zeros(logits.shape[0], dtype=torch.long).to(device)\n",
    "\n",
    "        logits = logits / temperature\n",
    "        return logits, labels"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bb5631261d0f6bd",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Training SimCLR\n",
    "\n",
    "Train the SimCLR model using the contrastive loss and augmented image pairs from CIFAR-10.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training with gpu: cuda.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 97/97 [00:05<00:00, 16.48it/s]\n",
      "100%|██████████| 97/97 [00:06<00:00, 15.77it/s]\n",
      "100%|██████████| 97/97 [00:06<00:00, 15.95it/s]\n",
      "100%|██████████| 97/97 [00:05<00:00, 16.44it/s]\n",
      "100%|██████████| 97/97 [00:05<00:00, 16.46it/s]\n",
      "100%|██████████| 97/97 [00:05<00:00, 16.50it/s]\n",
      "100%|██████████| 97/97 [00:05<00:00, 16.40it/s]\n",
      "100%|██████████| 97/97 [00:05<00:00, 16.25it/s]\n",
      "100%|██████████| 97/97 [00:05<00:00, 16.49it/s]\n",
      "100%|██████████| 97/97 [00:06<00:00, 15.82it/s]\n",
      "100%|██████████| 97/97 [00:05<00:00, 16.52it/s]\n",
      "100%|██████████| 97/97 [00:06<00:00, 15.56it/s]\n",
      "100%|██████████| 97/97 [00:05<00:00, 16.83it/s]\n",
      "100%|██████████| 97/97 [00:05<00:00, 16.59it/s]\n",
      "100%|██████████| 97/97 [00:05<00:00, 16.75it/s]\n",
      "100%|██████████| 97/97 [00:05<00:00, 16.45it/s]\n",
      "100%|██████████| 97/97 [00:06<00:00, 15.80it/s]\n",
      "100%|██████████| 97/97 [00:05<00:00, 16.67it/s]\n",
      "100%|██████████| 97/97 [00:06<00:00, 15.84it/s]\n",
      "100%|██████████| 97/97 [00:05<00:00, 16.36it/s]\n",
      "100%|██████████| 97/97 [00:06<00:00, 15.57it/s]\n",
      "100%|██████████| 97/97 [00:05<00:00, 16.38it/s]\n",
      " 54%|█████▎    | 52/97 [00:03<00:02, 18.11it/s]"
     ]
    }
   ],
   "source": [
    "from torch.utils.tensorboard import SummaryWriter\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "import logging\n",
    "from utils import accuracy, save_checkpoint\n",
    "\n",
    "\n",
    "with Timer():\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "    print(f\"Training with gpu: {device}.\")\n",
    "    # Initialize optimizer and loss criterion\n",
    "    model = ResNetSimCLR(base_model='resnet18', out_dim=128)\n",
    "    model = model.to(device)\n",
    "    lr = 3e-4\n",
    "    weight_decay = 1e-4\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr, weight_decay=weight_decay)\n",
    "    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=len(train_loader), eta_min=0,\n",
    "                                                           last_epoch=-1)\n",
    "    writer = SummaryWriter()\n",
    "    logging.basicConfig(filename=os.path.join(writer.log_dir, 'training.log'), level=logging.DEBUG)\n",
    "    criterion = torch.nn.CrossEntropyLoss().to(device)\n",
    "    # Set number of training epochs\n",
    "    epochs = 800\n",
    "    log_every_n_epochs = 1\n",
    "    logging.info(f\"Start SimCLR training for {epochs} epochs.\")\n",
    "    logging.info(f\"Training with gpu: {device}.\")\n",
    "    best_acc = 0\n",
    "    for epoch_counter in range(epochs):\n",
    "        loss_epoch = 0\n",
    "        for images, _ in tqdm(train_loader):\n",
    "            images = torch.cat(images, dim=0)\n",
    "\n",
    "            images = images.to(device)\n",
    "\n",
    "            # with autocast(enabled=fp16_precision):\n",
    "            features = model(images)\n",
    "            logits, labels = info_nce_loss(features)\n",
    "            loss = criterion(logits, labels)\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            loss_epoch += loss.item()\n",
    "            # scaler.scale(loss).backward()\n",
    "            # scaler.step(self.optimizer)\n",
    "            # scaler.update()\n",
    "        avg_loss = loss_epoch / len(train_loader)\n",
    "        # print(f\"Epoch {epoch_counter}:\\tLoss: {avg_loss}\")\n",
    "        # every log_every_n_epochs log epoch loss and accuracy\n",
    "        if epoch_counter % log_every_n_epochs == 0:\n",
    "            top1, top5 = accuracy(logits, labels, topk=(1, 5))\n",
    "            writer.add_scalar('loss', avg_loss, global_step=epoch_counter)\n",
    "            writer.add_scalar('acc/top1', top1[0], global_step=epoch_counter)\n",
    "            writer.add_scalar('acc/top5', top5[0], global_step=epoch_counter)\n",
    "            writer.add_scalar('learning_rate', scheduler.get_last_lr()[0], global_step=epoch_counter)\n",
    "            if top1[0] > best_acc:\n",
    "                best_acc = top1[0]\n",
    "                save_checkpoint({\n",
    "                    'epoch': epoch_counter,\n",
    "                    'arch': 'resnet18',\n",
    "                    'state_dict': model.state_dict(),\n",
    "                    'optimizer': optimizer.state_dict(),\n",
    "                }, is_best=True, filename=os.path.join(writer.log_dir, f'checkpoint_best.pth.tar'))\n",
    "\n",
    "\n",
    "        # warmup for the first 10 epochs\n",
    "        if epoch_counter >= 10:\n",
    "            scheduler.step()\n",
    "        logging.debug(f\"Epoch: {epoch_counter}\\tLoss: {loss}\\tTop1 accuracy: {top1[0]}\")\n",
    "\n",
    "    logging.info(\"Training has finished.\")\n",
    "    # save model checkpoints\n",
    "    checkpoint_name = 'checkpoint_{:04d}.pth.tar'.format(epochs)\n",
    "    save_checkpoint({\n",
    "        'epoch': epochs,\n",
    "        'arch': 'resnet18',\n",
    "        'state_dict': model.state_dict(),\n",
    "        'optimizer': optimizer.state_dict(),\n",
    "    }, is_best=False, filename=os.path.join(writer.log_dir, checkpoint_name))\n",
    "    logging.info(f\"Model checkpoint and metadata has been saved at {writer.log_dir}.\")"
   ],
   "metadata": {
    "collapsed": false,
    "is_executing": true,
    "ExecuteTime": {
     "start_time": "2023-09-26T15:15:26.715286939Z"
    }
   },
   "id": "5b7f76864c04709f"
  },
  {
   "cell_type": "markdown",
   "source": [
    "Load the model checkpoint and evaluate the learned representations using Linear Probing and KNN classification."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "e7a1fc915ff1ad65"
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = models.resnet18(pretrained=False, num_classes=10).to(device)\n",
    "# Load the checkpoint\n",
    "checkpoint_path = 'runs/Sep26_17-15-26_cpsadmin-Z790-AORUS-ELITE-AX/checkpoint_best.pth.tar'\n",
    "checkpoint = torch.load(checkpoint_path)\n",
    "state_dict = checkpoint['state_dict']\n",
    "# model.load_state_dict(state_dict)\n",
    "\n",
    "for k in list(state_dict.keys()):\n",
    "  if k.startswith('backbone.'):\n",
    "    if k.startswith('backbone') and not k.startswith('backbone.fc'):\n",
    "      # remove prefix\n",
    "      state_dict[k[len(\"backbone.\"):]] = state_dict[k]\n",
    "  del state_dict[k]\n",
    "log = model.load_state_dict(state_dict, strict=False)\n",
    "assert log.missing_keys == ['fc.weight', 'fc.bias']"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-10-06T12:33:04.099818694Z",
     "start_time": "2023-10-06T12:33:03.841207132Z"
    }
   },
   "id": "974bc2c895cb2485"
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "outputs": [],
   "source": [
    "# freeze all layers but the last fc\n",
    "for name, param in model.named_parameters():\n",
    "    if name not in ['fc.weight', 'fc.bias']:\n",
    "        param.requires_grad = False\n",
    "\n",
    "parameters = list(filter(lambda p: p.requires_grad, model.parameters()))\n",
    "assert len(parameters) == 2  # fc.weight, fc.bias"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-10-06T12:33:04.562825113Z",
     "start_time": "2023-10-06T12:33:04.554174461Z"
    }
   },
   "id": "93f72ecf304a489f"
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "outputs": [],
   "source": [
    "from torchvision import datasets\n",
    "def get_cifar10_data_loaders(download, shuffle=False, batch_size=256):\n",
    "  train_dataset = datasets.CIFAR10('./data', train=True, download=download,\n",
    "                                  transform=transforms.ToTensor())\n",
    "\n",
    "  train_loader = DataLoader(train_dataset, batch_size=batch_size,\n",
    "                            num_workers=0, drop_last=False, shuffle=shuffle)\n",
    "  \n",
    "  test_dataset = datasets.CIFAR10('./data', train=False, download=download,\n",
    "                                  transform=transforms.ToTensor())\n",
    "\n",
    "  test_loader = DataLoader(test_dataset, batch_size=2*batch_size,\n",
    "                            num_workers=10, drop_last=False, shuffle=shuffle)\n",
    "  return train_loader, test_loader\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-10-06T12:33:05.050502546Z",
     "start_time": "2023-10-06T12:33:05.043928699Z"
    }
   },
   "id": "6e880c61a3136e20"
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "optimizer = torch.optim.Adam(model.parameters(), lr=3e-4, weight_decay=0.0008)\n",
    "criterion = torch.nn.CrossEntropyLoss().to(device)\n",
    "train_loader, test_loader = get_cifar10_data_loaders(download=True)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-10-06T12:33:06.605497644Z",
     "start_time": "2023-10-06T12:33:05.400267070Z"
    }
   },
   "id": "a79946219b2fb147"
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0:\tTrain Accuracy: 61.83\tTest Accuracy: 69.19\tTest Top-5 Accuracy: 96.96\n",
      "Epoch 1:\tTrain Accuracy: 70.85\tTest Accuracy: 70.32\tTest Top-5 Accuracy: 97.47\n",
      "Epoch 2:\tTrain Accuracy: 71.92\tTest Accuracy: 70.97\tTest Top-5 Accuracy: 97.77\n",
      "Epoch 3:\tTrain Accuracy: 72.63\tTest Accuracy: 71.61\tTest Top-5 Accuracy: 97.91\n",
      "Epoch 4:\tTrain Accuracy: 73.19\tTest Accuracy: 72.12\tTest Top-5 Accuracy: 97.96\n",
      "Epoch 5:\tTrain Accuracy: 73.63\tTest Accuracy: 72.54\tTest Top-5 Accuracy: 98.13\n",
      "Epoch 6:\tTrain Accuracy: 73.96\tTest Accuracy: 72.88\tTest Top-5 Accuracy: 98.11\n",
      "Epoch 7:\tTrain Accuracy: 74.25\tTest Accuracy: 73.17\tTest Top-5 Accuracy: 98.17\n",
      "Epoch 8:\tTrain Accuracy: 74.49\tTest Accuracy: 73.42\tTest Top-5 Accuracy: 98.21\n",
      "Epoch 9:\tTrain Accuracy: 74.68\tTest Accuracy: 73.75\tTest Top-5 Accuracy: 98.25\n",
      "Elapsed time: 47.95 seconds\n"
     ]
    }
   ],
   "source": [
    "from utils import accuracy\n",
    "epochs = 10\n",
    "with Timer():\n",
    "    for epoch in range(epochs):\n",
    "        top1_train_accuracy = 0\n",
    "        model.train()\n",
    "        for counter, (x_batch, y_batch) in enumerate(train_loader):\n",
    "            x_batch = x_batch.to(device)\n",
    "            y_batch = y_batch.to(device)\n",
    "            \n",
    "            logits = model(x_batch)\n",
    "            loss = criterion(logits, y_batch)\n",
    "            top1 = accuracy(logits, y_batch, topk=(1,))\n",
    "            top1_train_accuracy += top1[0]\n",
    "            \n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            \n",
    "        top1_train_accuracy /= (counter + 1)\n",
    "        top1_accuracy = 0\n",
    "        top5_accuracy = 0\n",
    "        model.eval()\n",
    "        with torch.no_grad():\n",
    "            for counter, (x_batch, y_batch) in enumerate(test_loader):\n",
    "                x_batch = x_batch.to(device)\n",
    "                y_batch = y_batch.to(device)\n",
    "                \n",
    "                logits = model(x_batch)\n",
    "                \n",
    "                top1, top5 = accuracy(logits, y_batch, topk=(1,5))\n",
    "                top1_accuracy += top1[0]\n",
    "                top5_accuracy += top5[0]\n",
    "        \n",
    "        top1_accuracy /= (counter + 1)\n",
    "        top5_accuracy /= (counter + 1)\n",
    "        print(f\"Epoch {epoch}:\\tTrain Accuracy: {top1_train_accuracy.item():.2f}\\tTest Accuracy: {top1_accuracy.item():.2f}\\tTest Top-5 Accuracy: {top5_accuracy.item():.2f}\")\n",
    "  \n",
    "  "
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-10-06T12:33:54.564066447Z",
     "start_time": "2023-10-06T12:33:06.609095560Z"
    }
   },
   "id": "39c86d430700e100"
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Train a ResNet18 model from scratch on CIFAR-10 using the sane augmentation strategy as SimCLR  "
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "855fe8cc26f958af"
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "outputs": [],
   "source": [
    "from torchvision import datasets\n",
    "def get_cifar10_data_loaders(download, shuffle=False, batch_size=256):\n",
    "  train_dataset = datasets.CIFAR10('./data', train=True, download=download,\n",
    "                                  transform=transforms.ToTensor())\n",
    "\n",
    "  train_loader = DataLoader(train_dataset, batch_size=batch_size,\n",
    "                            num_workers=0, drop_last=False, shuffle=shuffle)\n",
    "  \n",
    "  test_dataset = datasets.CIFAR10('./data', train=False, download=download,\n",
    "                                  transform=transforms.ToTensor())\n",
    "\n",
    "  test_loader = DataLoader(test_dataset, batch_size=2*batch_size,\n",
    "                            num_workers=10, drop_last=False, shuffle=shuffle)\n",
    "  return train_loader, test_loader\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-09-27T10:24:24.614160376Z",
     "start_time": "2023-09-27T10:24:24.611873264Z"
    }
   },
   "id": "74586d7d4739fb2d"
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "from torchvision.models import resnet18\n",
    "model = resnet18(pretrained=False, num_classes=10).to(device)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=3e-4, weight_decay=0.0008)\n",
    "criterion = torch.nn.CrossEntropyLoss().to(device)\n",
    "train_loader, test_loader = get_cifar10_data_loaders(download=True)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-09-27T10:26:23.084698888Z",
     "start_time": "2023-09-27T10:26:22.191547635Z"
    }
   },
   "id": "8ab2d774394f9728"
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0:\tTrain Accuracy: 47.28\tTest Accuracy: 56.04\tTest Top-5 Accuracy: 95.10\n",
      "Epoch 1:\tTrain Accuracy: 62.14\tTest Accuracy: 61.73\tTest Top-5 Accuracy: 96.58\n",
      "Epoch 2:\tTrain Accuracy: 70.36\tTest Accuracy: 62.85\tTest Top-5 Accuracy: 96.75\n",
      "Epoch 3:\tTrain Accuracy: 76.94\tTest Accuracy: 63.65\tTest Top-5 Accuracy: 96.62\n",
      "Epoch 4:\tTrain Accuracy: 80.06\tTest Accuracy: 62.81\tTest Top-5 Accuracy: 96.05\n",
      "Epoch 5:\tTrain Accuracy: 82.35\tTest Accuracy: 64.39\tTest Top-5 Accuracy: 96.06\n",
      "Epoch 6:\tTrain Accuracy: 85.63\tTest Accuracy: 63.97\tTest Top-5 Accuracy: 95.74\n",
      "Epoch 7:\tTrain Accuracy: 88.24\tTest Accuracy: 64.20\tTest Top-5 Accuracy: 95.90\n",
      "Epoch 8:\tTrain Accuracy: 90.05\tTest Accuracy: 65.42\tTest Top-5 Accuracy: 96.06\n",
      "Epoch 9:\tTrain Accuracy: 91.83\tTest Accuracy: 65.22\tTest Top-5 Accuracy: 95.89\n",
      "Elapsed time: 26.27 seconds\n"
     ]
    }
   ],
   "source": [
    "from utils import accuracy\n",
    "epochs = 10\n",
    "with Timer():\n",
    "    for epoch in range(epochs):\n",
    "        top1_train_accuracy_sup = 0\n",
    "        for counter, (x_batch, y_batch) in enumerate(train_loader):\n",
    "            x_batch = x_batch.to(device)\n",
    "            y_batch = y_batch.to(device)\n",
    "            \n",
    "            logits = model(x_batch)\n",
    "            loss = criterion(logits, y_batch)\n",
    "            top1 = accuracy(logits, y_batch, topk=(1,))\n",
    "            top1_train_accuracy_sup += top1[0]\n",
    "            \n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            \n",
    "        top1_train_accuracy_sup /= (counter + 1)\n",
    "        top1_accuracy_sup = 0\n",
    "        top5_accuracy_sup = 0\n",
    "        for counter, (x_batch, y_batch) in enumerate(test_loader):\n",
    "            x_batch = x_batch.to(device)\n",
    "            y_batch = y_batch.to(device)\n",
    "            \n",
    "            logits = model(x_batch)\n",
    "            \n",
    "            top1, top5 = accuracy(logits, y_batch, topk=(1,5))\n",
    "            top1_accuracy_sup += top1[0]\n",
    "            top5_accuracy_sup += top5[0]\n",
    "        \n",
    "        top1_accuracy_sup /= (counter + 1)\n",
    "        top5_accuracy_sup /= (counter + 1)\n",
    "        print(f\"Epoch {epoch}:\\tTrain Accuracy: {top1_train_accuracy_sup.item():.2f}\\tTest Accuracy: {top1_accuracy_sup.item():.2f}\\tTest Top-5 Accuracy: {top5_accuracy_sup.item():.2f}\")\n",
    "            "
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-09-27T10:26:49.956507410Z",
     "start_time": "2023-09-27T10:26:23.682278251Z"
    }
   },
   "id": "9dcfa7b15c8395e7"
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "model = resnet18(pretrained=True).to(device)\n",
    "# overwrite the last fc layer\n",
    "model.fc = nn.Linear(512, 10).to(device)\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=3e-4, weight_decay=0.0008)\n",
    "criterion = torch.nn.CrossEntropyLoss().to(device)\n",
    "train_loader, test_loader = get_cifar10_data_loaders(download=True)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-09-27T10:26:50.840476452Z",
     "start_time": "2023-09-27T10:26:49.958244548Z"
    }
   },
   "id": "d3d787cb5d3d66dc"
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0:\tTrain Accuracy: 68.73\tTest Accuracy: 78.05\tTest Top-5 Accuracy: 98.70\n",
      "Epoch 1:\tTrain Accuracy: 83.82\tTest Accuracy: 80.25\tTest Top-5 Accuracy: 98.88\n",
      "Epoch 2:\tTrain Accuracy: 90.42\tTest Accuracy: 80.20\tTest Top-5 Accuracy: 98.74\n",
      "Epoch 3:\tTrain Accuracy: 92.74\tTest Accuracy: 80.11\tTest Top-5 Accuracy: 98.83\n",
      "Epoch 4:\tTrain Accuracy: 94.23\tTest Accuracy: 80.26\tTest Top-5 Accuracy: 98.75\n",
      "Epoch 5:\tTrain Accuracy: 95.63\tTest Accuracy: 80.89\tTest Top-5 Accuracy: 98.73\n",
      "Epoch 6:\tTrain Accuracy: 96.48\tTest Accuracy: 80.96\tTest Top-5 Accuracy: 98.85\n",
      "Epoch 7:\tTrain Accuracy: 97.06\tTest Accuracy: 81.44\tTest Top-5 Accuracy: 98.94\n",
      "Epoch 8:\tTrain Accuracy: 97.31\tTest Accuracy: 81.00\tTest Top-5 Accuracy: 98.46\n",
      "Epoch 9:\tTrain Accuracy: 97.51\tTest Accuracy: 81.58\tTest Top-5 Accuracy: 98.76\n",
      "Elapsed time: 26.06 seconds\n"
     ]
    }
   ],
   "source": [
    "from utils import accuracy\n",
    "epochs = 10\n",
    "with Timer():\n",
    "    for epoch in range(epochs):\n",
    "        top1_train_accuracy_sup_pre = 0\n",
    "        for counter, (x_batch, y_batch) in enumerate(train_loader):\n",
    "            x_batch = x_batch.to(device)\n",
    "            y_batch = y_batch.to(device)\n",
    "            \n",
    "            logits = model(x_batch)\n",
    "            loss = criterion(logits, y_batch)\n",
    "            top1 = accuracy(logits, y_batch, topk=(1,))\n",
    "            top1_train_accuracy_sup_pre += top1[0]\n",
    "            \n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            \n",
    "        top1_train_accuracy_sup_pre /= (counter + 1)\n",
    "        top1_accuracy_sup_pre = 0\n",
    "        top5_accuracy_sup_pre = 0\n",
    "        for counter, (x_batch, y_batch) in enumerate(test_loader):\n",
    "            x_batch = x_batch.to(device)\n",
    "            y_batch = y_batch.to(device)\n",
    "            \n",
    "            logits = model(x_batch)\n",
    "            \n",
    "            top1, top5 = accuracy(logits, y_batch, topk=(1,5))\n",
    "            top1_accuracy_sup_pre += top1[0]\n",
    "            top5_accuracy_sup_pre += top5[0]\n",
    "        \n",
    "        top1_accuracy_sup_pre /= (counter + 1)\n",
    "        top5_accuracy_sup_pre /= (counter + 1)\n",
    "        # print every 10 epochs\n",
    "        print(f\"Epoch {epoch}:\\tTrain Accuracy: {top1_train_accuracy_sup_pre.item():.2f}\\tTest Accuracy: {top1_accuracy_sup_pre.item():.2f}\\tTest Top-5 Accuracy: {top5_accuracy_sup_pre.item():.2f}\")\n",
    "            "
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-09-27T10:27:16.948482236Z",
     "start_time": "2023-09-27T10:26:50.840390430Z"
    }
   },
   "id": "ddac36785f993370"
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model                    Train Accuracy      Test Accuracy       Test Top-5 Accuracy \n",
      "SimCLR                   74.66               73.41               98.17               \n",
      "Supervised               91.83               65.22               95.89               \n",
      "Supervised Pretrained    97.51               81.58               98.76               \n"
     ]
    }
   ],
   "source": [
    "# print the results of the SimCLR model and the supervised model in a table format\n",
    "print(f\"{'Model':<25}{'Train Accuracy':<20}{'Test Accuracy':<20}{'Test Top-5 Accuracy':<20}\")\n",
    "print(f\"{'SimCLR':<25}{top1_train_accuracy.item():<20.2f}{top1_accuracy.item():<20.2f}{top5_accuracy.item():<20.2f}\")\n",
    "print(f\"{'Supervised':<25}{top1_train_accuracy_sup.item():<20.2f}{top1_accuracy_sup.item():<20.2f}{top5_accuracy_sup.item():<20.2f}\")\n",
    "print(f\"{'Supervised Pretrained':<25}{top1_train_accuracy_sup_pre.item():<20.2f}{top1_accuracy_sup_pre.item():<20.2f}{top5_accuracy_sup_pre.item():<20.2f}\")"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-09-27T10:27:30.971758924Z",
     "start_time": "2023-09-27T10:27:30.958172927Z"
    }
   },
   "id": "692dd08a956a4246"
  },
  {
   "cell_type": "markdown",
   "source": [],
   "metadata": {
    "collapsed": false
   },
   "id": "8676515f2c730790"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}