eda.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "import torch\n",
    "import pdb\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "from PIL import Image, ImageOps, ImageFilter\n",
    "from torch import nn, optim\n",
    "import torch\n",
    "\n",
    "from main import GaussianBlur, Solarization, off_diagonal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transform:\n",
    "    def __init__(self):\n",
    "        self.transform = transforms.Compose([\n",
    "            transforms.RandomResizedCrop(224, interpolation=Image.BICUBIC),\n",
    "            transforms.RandomHorizontalFlip(p=0.5),\n",
    "            transforms.RandomApply(\n",
    "                [transforms.ColorJitter(brightness=0.4, contrast=0.4,\n",
    "                                        saturation=0.2, hue=0.1)],\n",
    "                p=0.8\n",
    "            ),\n",
    "            transforms.RandomGrayscale(p=0.2),\n",
    "            GaussianBlur(p=1.0),\n",
    "            Solarization(p=0.0),\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "                                 std=[0.229, 0.224, 0.225])\n",
    "        ])\n",
    "        self.transform_prime = transforms.Compose([\n",
    "            transforms.RandomResizedCrop(224, interpolation=Image.BICUBIC),\n",
    "            transforms.RandomHorizontalFlip(p=0.5),\n",
    "            transforms.RandomApply(\n",
    "                [transforms.ColorJitter(brightness=0.4, contrast=0.4,\n",
    "                                        saturation=0.2, hue=0.1)],\n",
    "                p=0.8\n",
    "            ),\n",
    "            transforms.RandomGrayscale(p=0.2),\n",
    "            GaussianBlur(p=0.1),\n",
    "            Solarization(p=0.2),\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "                                 std=[0.229, 0.224, 0.225])\n",
    "        ])\n",
    "\n",
    "    def __call__(self, x):\n",
    "        y1 = self.transform(x)\n",
    "        y2 = self.transform_prime(x)\n",
    "        return y1, y2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BarlowTwins(nn.Module):\n",
    "    def __init__(self, args):\n",
    "        super().__init__()\n",
    "        self.args = args\n",
    "        self.backbone = torchvision.models.resnet50(zero_init_residual=True)\n",
    "        self.backbone.fc = nn.Identity()\n",
    "        # pdb.set_trace()\n",
    "        # projector\n",
    "        sizes = [2048] + list(map(int, args.projector.split('-')))\n",
    "        print(sizes)\n",
    "        layers = []\n",
    "        for i in range(len(sizes) - 2):\n",
    "            layers.append(nn.Linear(sizes[i], sizes[i + 1], bias=False))\n",
    "            layers.append(nn.BatchNorm1d(sizes[i + 1]))\n",
    "            layers.append(nn.ReLU(inplace=True))\n",
    "        layers.append(nn.Linear(sizes[-2], sizes[-1], bias=False))\n",
    "        self.projector = nn.Sequential(*layers)\n",
    "\n",
    "        # normalization layer for the representations z1 and z2\n",
    "        self.bn = nn.BatchNorm1d(sizes[-1], affine=False)\n",
    "\n",
    "    def forward(self, y1, y2):\n",
    "        z1 = self.projector(self.backbone(y1))\n",
    "        z2 = self.projector(self.backbone(y2))\n",
    "\n",
    "        # empirical cross-correlation matrix\n",
    "        c = self.bn(z1).T @ self.bn(z2)\n",
    "\n",
    "        # sum the cross-correlation matrix between all gpus\n",
    "        c.div_(self.args.batch_size)\n",
    "#         torch.distributed.all_reduce(c)\n",
    "\n",
    "        # use --scale-loss to multiply the loss by a constant factor\n",
    "        # see the Issues section of the readme\n",
    "        on_diag = torch.diagonal(c).add_(-1).pow_(2).sum().mul(self.args.scale_loss)\n",
    "        off_diag = off_diagonal(c).pow_(2).sum().mul(self.args.scale_loss)\n",
    "        loss = on_diag + self.args.lambd * off_diag\n",
    "        return loss,c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Args():\n",
    "    def __init__(self,lambd, batch_size, scale_loss, projector):\n",
    "        self.lambd = lambd\n",
    "        self.batch_size = batch_size\n",
    "        self.scale_loss = scale_loss\n",
    "        self.projector = projector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [],
   "source": [
    "args = Args(0.005,8,1/32,'64-64-64')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2048, 64, 64, 64]\n"
     ]
    }
   ],
   "source": [
    "model = BarlowTwins(args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [],
   "source": [
    "backbone = torchvision.models.resnet50(zero_init_residual=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Linear(in_features=2048, out_features=1000, bias=True)"
      ]
     },
     "execution_count": 86,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "backbone.fc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = torchvision.datasets.ImageFolder('/mounts/data/proj/jabbar/torchvision/imagenet/tiny-imagenet-200/',Transform())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "120000"
      ]
     },
     "execution_count": 72,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "metadata": {},
   "outputs": [],
   "source": [
    "loader = torch.utils.data.DataLoader(dataset,batch_size=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch = next(iter(loader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0, 0])"
      ]
     },
     "execution_count": 107,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2"
      ]
     },
     "execution_count": 108,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(batch[0][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 3, 224, 224])"
      ]
     },
     "execution_count": 109,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch[0][1].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 3, 224, 224])"
      ]
     },
     "execution_count": 110,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch[0][0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "metadata": {},
   "outputs": [],
   "source": [
    "(y1, y2), _ = batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 3, 224, 224])"
      ]
     },
     "execution_count": 112,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y1.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss,c = model(y1,y2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([64, 64])"
      ]
     },
     "execution_count": 114,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "c.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "metadata": {},
   "outputs": [],
   "source": [
    "emb = model.projector(model.backbone(y1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 64])"
      ]
     },
     "execution_count": 118,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "emb.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}