cs231n-deep-learning-for-computer-vision/2023/assignments/a3/additional-notes/003_simclr.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37161ae7-9be9-4a90-bc6e-3ffc05fb577a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tempfile\n",
    "from pathlib import Path\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import torch\n",
    "import torchvision\n",
    "from torchvision import transforms\n",
    "from torchvision.utils import save_image\n",
    "\n",
    "from utils import batch_plot, seed_everything\n",
    "\n",
    "\n",
    "np.set_printoptions(precision=2)\n",
    "pd.set_option(\"display.precision\", 2)\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "seed_everything()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1990d2cb-75c9-4f28-a389-8be95539cc7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 4\n",
    "image_size = 32\n",
    "feature_dims = 128\n",
    "temperature = 0.1\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "tmp_dir = Path(tempfile.mkdtemp())\n",
    "\n",
    "\n",
    "# https://github.com/sthalles/SimCLR\n",
    "class ContrastiveLearningViewGenerator(object):\n",
    "    def __init__(self, base_transform, n_views=2):\n",
    "        self.base_transform = base_transform\n",
    "        self.n_views = n_views\n",
    "\n",
    "    def __call__(self, x):\n",
    "        return [self.base_transform(x) for i in range(self.n_views)]\n",
    "\n",
    "\n",
    "null_transforms = transforms.Compose(\n",
    "    [\n",
    "        transforms.ToTensor(),\n",
    "    ]\n",
    ")\n",
    "\n",
    "contrast_transforms = transforms.Compose(\n",
    "    [\n",
    "        transforms.RandomResizedCrop(size=image_size),\n",
    "        transforms.RandomHorizontalFlip(p=0.5),\n",
    "        transforms.RandomApply([transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1)], p=0.8),\n",
    "        transforms.RandomGrayscale(p=0.2),\n",
    "        transforms.GaussianBlur(kernel_size=int(0.1 * image_size)),\n",
    "        transforms.ToTensor(),\n",
    "    ]\n",
    ")\n",
    "\n",
    "\n",
    "original_dataset = torchvision.datasets.CIFAR10(\n",
    "    root=\"../code/cs231n/datasets\", train=True, download=True, transform=null_transforms\n",
    ")\n",
    "original_loader = torch.utils.data.DataLoader(original_dataset, batch_size=batch_size, shuffle=False)\n",
    "original_images, original_labels = next(iter(original_loader))\n",
    "batch_plot(xs=original_images.numpy().transpose((0, 2, 3, 1)), ys=original_labels.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "713d1193-e45d-4f8d-8ace-2091ab039d4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, image in enumerate(original_images):\n",
    "    save_image(image, tmp_dir.joinpath(f\"original_{i:02d}.png\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a514fb44-e8a8-4aed-87b3-7f5b15cc4a19",
   "metadata": {},
   "outputs": [],
   "source": [
    "contrast_dataset = torchvision.datasets.CIFAR10(\n",
    "    root=\"../code/cs231n/datasets\",\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=ContrastiveLearningViewGenerator(base_transform=contrast_transforms),\n",
    ")\n",
    "contrast_loader = torch.utils.data.DataLoader(contrast_dataset, batch_size=batch_size, shuffle=False)\n",
    "contrast_images, _ = next(iter(contrast_loader))\n",
    "contrast_images = torch.cat(contrast_images, dim=0)\n",
    "batch_plot(xs=contrast_images.numpy().transpose((0, 2, 3, 1)), rows=2, cols=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f2a696e-6c16-4bc7-beda-6963dd2db822",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, image in enumerate(contrast_images):\n",
    "    save_image(image, tmp_dir.joinpath(f\"constract_{i:02d}.png\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47fe9a2d-4b16-4eca-987b-e799820410c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "import torchvision.models as models\n",
    "\n",
    "\n",
    "class ResNetSimCLR(nn.Module):\n",
    "    # https://github.com/sthalles/SimCLR/blob/master/models/resnet_simclr.py\n",
    "    def __init__(self, base_model=\"resnet18\", **kwargs):\n",
    "        super().__init__()\n",
    "        # f(.)\n",
    "        self.encoder = models.__dict__[base_model](**kwargs)\n",
    "        dim_mlp = self.encoder.fc.in_features\n",
    "        # g(.)\n",
    "        self.encoder.fc = nn.Sequential(nn.Linear(dim_mlp, dim_mlp), nn.ReLU(), self.encoder.fc)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.encoder(x)\n",
    "\n",
    "\n",
    "def compute_sim_matrix_direct(features):\n",
    "    return torch.nn.functional.cosine_similarity(features[:, None, :], features[None, :, :], dim=-1)\n",
    "\n",
    "\n",
    "def compute_sim_matrix_functional(features):\n",
    "    features = torch.nn.functional.normalize(features, dim=1)\n",
    "    return features @ features.T\n",
    "\n",
    "\n",
    "def compute_sim_matrix_norm(features):\n",
    "    features = features / torch.linalg.norm(features, dim=1, keepdims=True)\n",
    "    return features @ features.T\n",
    "\n",
    "\n",
    "def simclr_loss_direct(features, temperature, n_views=2):\n",
    "    n = len(features) // n_views\n",
    "    similarity_matrix = compute_sim_matrix_norm(features)  # (2*N, 2*N))\n",
    "    exponential = (similarity_matrix / temperature).exp()  # [2*N, 2*N]\n",
    "    mask = torch.eye(2 * n, dtype=torch.bool)  # [2*N, 2*N]\n",
    "    denom = exponential[~mask].view(2 * n, -1).sum(axis=1)  # [2*N, 1]\n",
    "    loss = -(exponential / denom).log()\n",
    "    return loss[mask.roll(shifts=n, dims=0)].mean()\n",
    "\n",
    "\n",
    "def simclr_loss_separate(features, temperature, n_views=2):\n",
    "    n = len(features) // n_views\n",
    "    similarity_matrix = compute_sim_matrix_norm(features) / temperature  # (2*N, 2*N))\n",
    "    mask = torch.eye(2 * n, dtype=torch.bool)  # [2*N, 2*N]\n",
    "    similarity_matrix.masked_fill_(mask, -float(\"inf\"))  # [2*N, 2*N]\n",
    "    # -log(exp(a)/sum(exp(b))) = -a + logsumexp(b)\n",
    "    loss = -similarity_matrix[mask.roll(shifts=n, dims=0)] + torch.logsumexp(similarity_matrix, dim=-1)\n",
    "    return loss.mean()\n",
    "\n",
    "\n",
    "def simclr_loss_criterion(features, temperature, n_views=2):\n",
    "    n = len(features) // n_views\n",
    "    similarity_matrix = compute_sim_matrix_norm(features)  # (2*N, 2*N))\n",
    "    mask = torch.eye(2 * n, dtype=torch.bool)  # [2*N, 2*N]\n",
    "    similarity_matrix = similarity_matrix[~mask].view(2 * n, -1)\n",
    "    labels = mask.roll(shifts=n, dims=0)[~mask].view(2 * n, -1)\n",
    "    positives = similarity_matrix[labels].view(2 * n, -1)\n",
    "    negatives = similarity_matrix[~labels].view(2 * n, -1)\n",
    "    logits = torch.cat([positives, negatives], dim=1) / temperature\n",
    "    labels = torch.zeros(logits.shape[0], dtype=torch.long)\n",
    "    return torch.nn.functional.cross_entropy(logits, labels)\n",
    "\n",
    "\n",
    "model = ResNetSimCLR(weights=None, num_classes=feature_dims)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b15dbaf4-a3d3-4553-b619-8eeb3b81b6fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "features = model(contrast_images)\n",
    "\n",
    "assert features.shape == (2 * batch_size, feature_dims)\n",
    "\n",
    "\n",
    "torch.testing.assert_close(\n",
    "    compute_sim_matrix_direct(features=features), compute_sim_matrix_functional(features=features)\n",
    ")\n",
    "\n",
    "torch.testing.assert_close(compute_sim_matrix_norm(features=features), compute_sim_matrix_functional(features=features))\n",
    "\n",
    "torch.testing.assert_close(\n",
    "    simclr_loss_direct(features=features, temperature=temperature),\n",
    "    simclr_loss_criterion(features=features, temperature=temperature),\n",
    ")\n",
    "\n",
    "torch.testing.assert_close(\n",
    "    simclr_loss_direct(features=features, temperature=temperature),\n",
    "    simclr_loss_separate(features=features, temperature=temperature),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eeb39878-9134-44ee-b104-1513381fdb45",
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit compute_sim_matrix_norm(features=features)\n",
    "%timeit compute_sim_matrix_functional(features=features)\n",
    "%timeit compute_sim_matrix_direct(features=features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "918fb531-73f9-4907-90e4-2820d2229d12",
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit simclr_loss_direct(features=features, temperature=temperature)\n",
    "%timeit simclr_loss_criterion(features=features, temperature=temperature)\n",
    "%timeit simclr_loss_separate(features=features, temperature=temperature)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "!open {tmp_dir}"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "2839699d5906fe18"
  },
  {
   "cell_type": "markdown",
   "id": "55f0313f-b6f6-408a-983e-850fb754e5bc",
   "metadata": {},
   "source": [
    "# references\n",
    "\n",
    "- [Contrastive Representation Learning](https://lilianweng.github.io/posts/2021-05-31-contrastive/)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}