essential_notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Self-Supervised Relational Reasoning\n",
    "------------------------------------------------------\n",
    "\n",
    "Official PyTorch implementation of the paper: \n",
    "\n",
    "```\n",
    "\"Sefl-Supervised Relational Reasoning for Representation Learning\", Patacchiola, M., and Storkey, A., *Advances in Neural Information Processing Systems (NeurIPS)\n",
    "```\n",
    "\n",
    "In this notebook is presented an essential implementation of the method, which is modular and easy to extend. The code has been successfully tested on `Ubuntu 18.04 LTS` with `PyTorch 1.4`, `Torchvision 0.5`, and `PIL 7.0`, so we suggest a similar configuration. First of all we import all the necessary modules and print their versions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.4.0\n",
      "0.5.0\n",
      "7.0.0\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "from PIL import Image\n",
    "import math\n",
    "\n",
    "print(torch.__version__)\n",
    "print(torchvision.__version__)\n",
    "print(Image.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We override the CIFAR10 and CIFAR100 classes in torchvision, with a version able to return **multiple augmentations** of the input mini-batch. This can be easily done by overriding the `__getitem__()` method. We add the parameter `K` which specifies the number of augmentations we want to apply to the input image. The output of `__getitem__()` is a list of lenght `K` containing different augmentations of the input instance, and the label of that image (this is discarded during the unsupervised training phase)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiCIFAR10(torchvision.datasets.CIFAR10):\n",
    "  \"\"\"Override torchvision CIFAR10 for multi-image management.\n",
    "  Similar class can be defined for other datasets (e.g. CIFAR100).\n",
    "  Given K total augmentations, it returns a list of lenght K with\n",
    "  different augmentations of the input mini-batch.\n",
    "  \"\"\"\n",
    "  def __init__(self, K, **kwds):\n",
    "    super().__init__(**kwds)\n",
    "    self.K = K # tot number of augmentations\n",
    "            \n",
    "  def __getitem__(self, index):\n",
    "    img, target = self.data[index], self.targets[index]\n",
    "    pic = Image.fromarray(img)            \n",
    "    img_list = list()\n",
    "    if self.transform is not None:\n",
    "      for _ in range(self.K):\n",
    "        img_transformed = self.transform(pic.copy())\n",
    "        img_list.append(img_transformed)\n",
    "    else:\n",
    "        img_list = img\n",
    "    return img_list, target"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiCIFAR100(torchvision.datasets.CIFAR100):\n",
    "  \"\"\"Override torchvision CIFAR100 for multi-image management.\n",
    "  Given K total augmentations, it returns a list of lenght K with\n",
    "  different augmentations of the input mini-batch.\n",
    "  \"\"\"\n",
    "  def __init__(self, K, **kwds):\n",
    "    super().__init__(**kwds)\n",
    "    self.K = K # tot number of augmentations\n",
    "            \n",
    "  def __getitem__(self, index):\n",
    "    img, target = self.data[index], self.targets[index]\n",
    "    pic = Image.fromarray(img)            \n",
    "    img_list = list()\n",
    "    if self.transform is not None:\n",
    "      for _ in range(self.K):\n",
    "        img_transformed = self.transform(pic.copy())\n",
    "        img_list.append(img_transformed)\n",
    "    else:\n",
    "        img_list = img\n",
    "    return img_list, target"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we need a convolutional neural network (CNN) backbone, which is used as a preliminary stage for dimensionality reduction of the input images. Here, we define a simple 4-layers CNN, but any other network with a linear output layer can be used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Conv4(torch.nn.Module):\n",
    "    \"\"\"A simple 4 layers CNN.\n",
    "    Used as backbone.    \n",
    "    \"\"\"\n",
    "    def __init__(self):\n",
    "        super(Conv4, self).__init__()\n",
    "        self.feature_size = 64\n",
    "        self.name = \"conv4\"\n",
    "\n",
    "        self.layer1 = torch.nn.Sequential(\n",
    "          torch.nn.Conv2d(3, 8, kernel_size=3, stride=1, padding=1, bias=False),\n",
    "          torch.nn.BatchNorm2d(8),\n",
    "          torch.nn.ReLU(),\n",
    "          torch.nn.AvgPool2d(kernel_size=2, stride=2)\n",
    "        )\n",
    "\n",
    "        self.layer2 = torch.nn.Sequential(\n",
    "          torch.nn.Conv2d(8, 16, kernel_size=3, stride=1, padding=1, bias=False),\n",
    "          torch.nn.BatchNorm2d(16),\n",
    "          torch.nn.ReLU(),\n",
    "          torch.nn.AvgPool2d(kernel_size=2, stride=2)\n",
    "        )\n",
    "\n",
    "        self.layer3 = torch.nn.Sequential(\n",
    "          torch.nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1, bias=False),\n",
    "          torch.nn.BatchNorm2d(32),\n",
    "          torch.nn.ReLU(),\n",
    "          torch.nn.AvgPool2d(kernel_size=2, stride=2)\n",
    "        )\n",
    "\n",
    "        self.layer4 = torch.nn.Sequential(\n",
    "          torch.nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1, bias=False),\n",
    "          torch.nn.BatchNorm2d(64),\n",
    "          torch.nn.ReLU(),\n",
    "          torch.nn.AdaptiveAvgPool2d(1)\n",
    "        )\n",
    "\n",
    "        self.flatten = torch.nn.Flatten()\n",
    "\n",
    "        for m in self.modules():\n",
    "            if isinstance(m, torch.nn.Conv2d):\n",
    "                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels\n",
    "                m.weight.data.normal_(0, math.sqrt(2. / n))\n",
    "            elif isinstance(m, torch.nn.BatchNorm2d):\n",
    "                m.weight.data.fill_(1)\n",
    "                m.bias.data.zero_()\n",
    "\n",
    "    def forward(self, x):\n",
    "        h = self.layer1(x)\n",
    "        h = self.layer2(h)\n",
    "        h = self.layer3(h)\n",
    "        h = self.layer4(h)\n",
    "        h = self.flatten(h)\n",
    "        return h"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we define the Relational Reasoning class. This consists of an `__inti__()` method, an internal `aggregate()` method, and finally a `train()` routine. \n",
    "\n",
    "In the `__init__()` we pass the CNN backbone, and the feature size representing the number of output (linear) units in the CNN. At init time, the relation head is created. This is just a multi-layer perceptron (MLP) with 256 hidden units and leaky-ReLU activation function. The type of relation head is important. If the relation head is too complex then it can easily discriminate the relation pairs; as a result the backbone will not learn useful representations.\n",
    "\n",
    "The `aggregate()` method takes as input the features produced in the forward pass by the backbone and `K` which is the total number of augmentations we are using. The output of the aggregation phase are the relation pairs joined by a concatenation operator (other commutative operators can be used, but concatenation is the most effective). \n",
    "\n",
    "The `train()` routine is just an iterative learning schedule for optimizing the parameters of the backbone and relation head. It takes as input a train loader, and an integer representing the total number of epochs. Here we use the Binary Cross-Entropy loss (BCE) but using a Focal Loss can give some boost."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RelationalReasoning(torch.nn.Module):\n",
    "  \"\"\"Self-Supervised Relational Reasoning.\n",
    "  Essential implementation of the method, which uses\n",
    "  the 'cat' aggregation function (the most effective),\n",
    "  and can be used with any backbone.\n",
    "  \"\"\"\n",
    "  def __init__(self, backbone, feature_size=64):\n",
    "    super(RelationalReasoning, self).__init__()\n",
    "    self.backbone = backbone\n",
    "    self.relation_head = torch.nn.Sequential(\n",
    "                             torch.nn.Linear(feature_size*2, 256),\n",
    "                             torch.nn.BatchNorm1d(256),\n",
    "                             torch.nn.LeakyReLU(),\n",
    "                             torch.nn.Linear(256, 1))\n",
    "\n",
    "  def aggregate(self, features, K):\n",
    "    relation_pairs_list = list()\n",
    "    targets_list = list()\n",
    "    size = int(features.shape[0] / K)\n",
    "    shifts_counter=1\n",
    "    for index_1 in range(0, size*K, size):\n",
    "      for index_2 in range(index_1+size, size*K, size):\n",
    "        # Using the 'cat' aggregation function by default\n",
    "        pos_pair = torch.cat([features[index_1:index_1+size], \n",
    "                              features[index_2:index_2+size]], 1)\n",
    "        # Shuffle without collisions by rolling the mini-batch (negatives)\n",
    "        neg_pair = torch.cat([\n",
    "                     features[index_1:index_1+size], \n",
    "                     torch.roll(features[index_2:index_2+size], \n",
    "                     shifts=shifts_counter, dims=0)], 1)\n",
    "        relation_pairs_list.append(pos_pair)\n",
    "        relation_pairs_list.append(neg_pair)\n",
    "        targets_list.append(torch.ones(size, dtype=torch.float32))\n",
    "        targets_list.append(torch.zeros(size, dtype=torch.float32))\n",
    "        shifts_counter+=1\n",
    "        if(shifts_counter>=size): \n",
    "            shifts_counter=1 # avoid identity pairs\n",
    "    relation_pairs = torch.cat(relation_pairs_list, 0)\n",
    "    targets = torch.cat(targets_list, 0)\n",
    "    return relation_pairs, targets\n",
    "\n",
    "  def train(self, tot_epochs, train_loader):\n",
    "    optimizer = torch.optim.Adam([\n",
    "                  {'params': self.backbone.parameters()},\n",
    "                  {'params': self.relation_head.parameters()}])                               \n",
    "    BCE = torch.nn.BCEWithLogitsLoss()\n",
    "    self.backbone.train()\n",
    "    self.relation_head.train()\n",
    "    for epoch in range(tot_epochs):\n",
    "      # the real target is discarded (unsupervised)\n",
    "      for i, (data_augmented, _) in enumerate(train_loader):\n",
    "        K = len(data_augmented) # tot augmentations\n",
    "        x = torch.cat(data_augmented, 0)\n",
    "        optimizer.zero_grad()              \n",
    "        # forward pass (backbone)\n",
    "        features = self.backbone(x) \n",
    "        # aggregation function\n",
    "        relation_pairs, targets = self.aggregate(features, K)\n",
    "        # forward pass (relation head)\n",
    "        score = self.relation_head(relation_pairs).squeeze()        \n",
    "        # cross-entropy loss and backward\n",
    "        loss = BCE(score, targets)\n",
    "        loss.backward()\n",
    "        optimizer.step()            \n",
    "        # estimate the accuracy\n",
    "        predicted = torch.round(torch.sigmoid(score))\n",
    "        correct = predicted.eq(targets.view_as(predicted)).sum()\n",
    "        accuracy = (100.0 * correct / float(len(targets)))\n",
    "        \n",
    "        if(i%100==0):\n",
    "          print('Epoch [{}][{}/{}] loss: {:.5f}; accuracy: {:.2f}%' \\\n",
    "            .format(epoch+1, i+1, len(train_loader)+1, \n",
    "                    loss.item(), accuracy.item()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unsupervised training\n",
    "--------------------------------\n",
    "\n",
    "In this section we use self-supervised relational reasoning for training a Conv-4 backbone on the unsupervised CIFAR-10 dataset.\n",
    "\n",
    "In the next cell we define some hyper-parameters, such as `K` the number of total augmentations, the mini-batch size, tot_epochs and feature size (related to the backbone used). The time complexity is quadratic in the number of augmentations `K`, therefore here we use a small value just for checking the code. In the paper we used `K=32` for CIFAR-10 and CIFAR-100 experiments with mini-batch 64. In this example we train the relational model for just *10 epochs*, in the paper we trained for 200 epochs. We invite the reader to experiment with the hyper-parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hyper-parameters of the simulation\n",
    "K = 4 # tot augmentations, in the paper K=32 for CIFAR10/100\n",
    "batch_size = 64 # 64 has been used in the paper\n",
    "tot_epochs = 10 # 200 has been used in the paper\n",
    "feature_size = 64 # number of units for the Conv4 backbone"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below we define the augmentation strategy. Note that here we use the CIFAR-10 normalization values, which must be changed if CIFAR-100 is used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Those are the transformations used in the paper\n",
    "normalize = transforms.Normalize(mean=[0.491, 0.482, 0.447], \n",
    "                                 std=[0.247, 0.243, 0.262]) # CIFAR10\n",
    "#normalize = transforms.Normalize(mean=[0.507, 0.487, 0.441], \n",
    "#                                 std=[0.267, 0.256, 0.276]) # CIFAR100\n",
    "\n",
    "color_jitter = transforms.ColorJitter(brightness=0.8, contrast=0.8, \n",
    "                                      saturation=0.8, hue=0.2)\n",
    "rnd_color_jitter = transforms.RandomApply([color_jitter], p=0.8)\n",
    "rnd_gray = transforms.RandomGrayscale(p=0.2)\n",
    "rnd_rcrop = transforms.RandomResizedCrop(size=32, scale=(0.08, 1.0), \n",
    "                                         interpolation=2)\n",
    "rnd_hflip = transforms.RandomHorizontalFlip(p=0.5)\n",
    "train_transform = transforms.Compose([rnd_rcrop, rnd_hflip,\n",
    "                                      rnd_color_jitter, rnd_gray, \n",
    "                                      transforms.ToTensor(), normalize])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Defining the backbone, model, and train loader. This will download the dataset if it is not available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "backbone = Conv4() # simple CNN with 64 linear output units\n",
    "model = RelationalReasoning(backbone, feature_size)    \n",
    "train_set = MultiCIFAR10(K=K, root='data', train=True, \n",
    "                         transform=train_transform, \n",
    "                         download=True)\n",
    "train_loader = torch.utils.data.DataLoader(train_set, \n",
    "                                           batch_size=batch_size, \n",
    "                                           shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next cell starts the training for the number of epochs specified and then save the backbone in a local file. This may take a while, based on your hardware configuration, the size of `K`, and the number of epochs. The code can be easily adapted to run on a GPU if you have one. The following cell takes ~30 minutes to complete the 10 training epochs on a medium-level laptop (with no GPU acceleration). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1][1/783] loss: 0.70772; accuracy: 51.56%\n",
      "Epoch [1][101/783] loss: 0.63008; accuracy: 62.50%\n",
      "Epoch [1][201/783] loss: 0.58195; accuracy: 69.53%\n",
      "Epoch [1][301/783] loss: 0.58450; accuracy: 67.97%\n",
      "Epoch [1][401/783] loss: 0.60491; accuracy: 65.23%\n",
      "Epoch [1][501/783] loss: 0.53203; accuracy: 72.92%\n",
      "Epoch [1][601/783] loss: 0.57815; accuracy: 68.62%\n",
      "Epoch [1][701/783] loss: 0.52331; accuracy: 73.18%\n",
      "Epoch [2][1/783] loss: 0.49492; accuracy: 74.22%\n",
      "Epoch [2][101/783] loss: 0.48556; accuracy: 77.08%\n",
      "Epoch [2][201/783] loss: 0.44053; accuracy: 79.17%\n",
      "Epoch [2][301/783] loss: 0.52381; accuracy: 71.88%\n",
      "Epoch [2][401/783] loss: 0.53081; accuracy: 73.05%\n",
      "Epoch [2][501/783] loss: 0.44579; accuracy: 77.34%\n",
      "Epoch [2][601/783] loss: 0.48443; accuracy: 76.04%\n",
      "Epoch [2][701/783] loss: 0.49158; accuracy: 73.44%\n",
      "Epoch [3][1/783] loss: 0.43537; accuracy: 79.82%\n",
      "Epoch [3][101/783] loss: 0.45528; accuracy: 78.65%\n",
      "Epoch [3][201/783] loss: 0.46668; accuracy: 78.52%\n",
      "Epoch [3][301/783] loss: 0.42457; accuracy: 77.73%\n",
      "Epoch [3][401/783] loss: 0.49923; accuracy: 76.04%\n",
      "Epoch [3][501/783] loss: 0.45702; accuracy: 77.73%\n",
      "Epoch [3][601/783] loss: 0.43920; accuracy: 78.65%\n",
      "Epoch [3][701/783] loss: 0.46995; accuracy: 76.56%\n",
      "Epoch [4][1/783] loss: 0.46817; accuracy: 76.69%\n",
      "Epoch [4][101/783] loss: 0.40496; accuracy: 82.68%\n",
      "Epoch [4][201/783] loss: 0.47275; accuracy: 76.43%\n",
      "Epoch [4][301/783] loss: 0.36549; accuracy: 82.55%\n",
      "Epoch [4][401/783] loss: 0.40239; accuracy: 82.29%\n",
      "Epoch [4][501/783] loss: 0.42898; accuracy: 80.99%\n",
      "Epoch [4][601/783] loss: 0.34945; accuracy: 83.33%\n",
      "Epoch [4][701/783] loss: 0.38465; accuracy: 82.29%\n",
      "Epoch [5][1/783] loss: 0.39593; accuracy: 82.16%\n",
      "Epoch [5][101/783] loss: 0.42994; accuracy: 80.08%\n",
      "Epoch [5][201/783] loss: 0.41548; accuracy: 80.47%\n",
      "Epoch [5][301/783] loss: 0.36508; accuracy: 83.20%\n",
      "Epoch [5][401/783] loss: 0.42240; accuracy: 81.38%\n",
      "Epoch [5][501/783] loss: 0.36961; accuracy: 83.33%\n",
      "Epoch [5][601/783] loss: 0.36313; accuracy: 83.46%\n",
      "Epoch [5][701/783] loss: 0.37241; accuracy: 82.68%\n",
      "Epoch [6][1/783] loss: 0.35195; accuracy: 83.72%\n",
      "Epoch [6][101/783] loss: 0.40442; accuracy: 81.51%\n",
      "Epoch [6][201/783] loss: 0.35945; accuracy: 83.72%\n",
      "Epoch [6][301/783] loss: 0.36455; accuracy: 82.68%\n",
      "Epoch [6][401/783] loss: 0.38897; accuracy: 82.29%\n",
      "Epoch [6][501/783] loss: 0.36484; accuracy: 82.16%\n",
      "Epoch [6][601/783] loss: 0.43785; accuracy: 79.95%\n",
      "Epoch [6][701/783] loss: 0.40670; accuracy: 81.77%\n",
      "Epoch [7][1/783] loss: 0.34668; accuracy: 83.59%\n",
      "Epoch [7][101/783] loss: 0.36399; accuracy: 84.38%\n",
      "Epoch [7][201/783] loss: 0.35402; accuracy: 83.98%\n",
      "Epoch [7][301/783] loss: 0.33216; accuracy: 85.94%\n",
      "Epoch [7][401/783] loss: 0.33434; accuracy: 84.90%\n",
      "Epoch [7][501/783] loss: 0.35985; accuracy: 83.20%\n",
      "Epoch [7][601/783] loss: 0.33478; accuracy: 86.20%\n",
      "Epoch [7][701/783] loss: 0.33205; accuracy: 85.94%\n",
      "Epoch [8][1/783] loss: 0.31099; accuracy: 86.07%\n",
      "Epoch [8][101/783] loss: 0.38285; accuracy: 83.33%\n",
      "Epoch [8][201/783] loss: 0.35258; accuracy: 83.20%\n",
      "Epoch [8][301/783] loss: 0.32939; accuracy: 84.64%\n",
      "Epoch [8][401/783] loss: 0.35715; accuracy: 83.46%\n",
      "Epoch [8][501/783] loss: 0.36472; accuracy: 83.33%\n",
      "Epoch [8][601/783] loss: 0.35551; accuracy: 84.64%\n",
      "Epoch [8][701/783] loss: 0.35802; accuracy: 83.33%\n",
      "Epoch [9][1/783] loss: 0.37295; accuracy: 82.68%\n",
      "Epoch [9][101/783] loss: 0.33368; accuracy: 85.03%\n",
      "Epoch [9][201/783] loss: 0.34582; accuracy: 84.11%\n",
      "Epoch [9][301/783] loss: 0.33333; accuracy: 84.24%\n",
      "Epoch [9][401/783] loss: 0.34445; accuracy: 84.51%\n",
      "Epoch [9][501/783] loss: 0.35353; accuracy: 83.20%\n",
      "Epoch [9][601/783] loss: 0.30412; accuracy: 85.81%\n",
      "Epoch [9][701/783] loss: 0.36279; accuracy: 84.64%\n",
      "Epoch [10][1/783] loss: 0.29136; accuracy: 88.15%\n",
      "Epoch [10][101/783] loss: 0.28832; accuracy: 87.63%\n",
      "Epoch [10][201/783] loss: 0.29915; accuracy: 86.33%\n",
      "Epoch [10][301/783] loss: 0.30783; accuracy: 87.11%\n",
      "Epoch [10][401/783] loss: 0.34802; accuracy: 83.72%\n",
      "Epoch [10][501/783] loss: 0.34900; accuracy: 83.33%\n",
      "Epoch [10][601/783] loss: 0.29620; accuracy: 86.07%\n",
      "Epoch [10][701/783] loss: 0.34308; accuracy: 85.29%\n"
     ]
    }
   ],
   "source": [
    "model.train(tot_epochs=tot_epochs, train_loader=train_loader)\n",
    "torch.save(model.backbone.state_dict(), './backbone.tar')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Linear evaluation\n",
    "------------------------\n",
    "\n",
    "Once the model has been trained and the backbone saved, we can use the backbone for downstream tasks such as classification or regression. Here, we perform a **linear evaluation** test which takes the backbone, stack a linear layer on top of it, then train just the weights of the linear classifier (no backprop on the backbone). This allow us to check the quality of the representations, and how close they are to the fully supervised upper-bound score. We perform linear evaluation on the same dataset (e.g. CIFAR-10) by accessing the labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# no augmentations used for linear evaluation\n",
    "transform_lineval = transforms.Compose([transforms.ToTensor(), normalize]) \n",
    "train_set_lineval = torchvision.datasets.CIFAR10('data', train=True, transform=transform_lineval)\n",
    "test_set_lineval = torchvision.datasets.CIFAR10('data', train=False, transform=transform_lineval)\n",
    "train_loader_lineval = torch.utils.data.DataLoader(train_set_lineval, batch_size=128, shuffle=True)\n",
    "test_loader_lineval = torch.utils.data.DataLoader(test_set_lineval, batch_size=128, shuffle=False)\n",
    "# 64 are the number of output features in the backbone, and 10 the number of classes\n",
    "linear_layer = torch.nn.Linear(64, 10)\n",
    "# loading the saved backbone\n",
    "backbone_lineval = Conv4() #defining a raw backbone model\n",
    "checkpoint = torch.load('./backbone.tar')\n",
    "backbone_lineval.load_state_dict(checkpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we start the training routine on the supervised CIFAR-10 for 10 epochs. This phase is much faster because we are just backpropagating on the linear layer. This cell takes ~2 minutes to complete the 10 epochs (medium-level laptop with no GPU acceleration)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Linear evaluation\n",
      "Epoch [1] loss: 1.47552; accuracy: 52.99%\n",
      "Epoch [2] loss: 1.22791; accuracy: 52.98%\n",
      "Epoch [3] loss: 1.44275; accuracy: 53.03%\n",
      "Epoch [4] loss: 1.32339; accuracy: 53.02%\n",
      "Epoch [5] loss: 1.47076; accuracy: 53.05%\n",
      "Epoch [6] loss: 1.33488; accuracy: 53.07%\n",
      "Epoch [7] loss: 1.22281; accuracy: 53.15%\n",
      "Epoch [8] loss: 1.29908; accuracy: 53.11%\n",
      "Epoch [9] loss: 1.34019; accuracy: 53.12%\n",
      "Epoch [10] loss: 1.29688; accuracy: 53.16%\n"
     ]
    }
   ],
   "source": [
    "optimizer = torch.optim.Adam(linear_layer.parameters())                               \n",
    "CE = torch.nn.CrossEntropyLoss()\n",
    "linear_layer.train()\n",
    "backbone_lineval.eval()\n",
    "\n",
    "print('Linear evaluation')\n",
    "for epoch in range(10):\n",
    "    accuracy_list = list()\n",
    "    for i, (data, target) in enumerate(train_loader_lineval):\n",
    "        optimizer.zero_grad()\n",
    "        output = backbone_lineval(data).detach()\n",
    "        output = linear_layer(output)\n",
    "        loss = CE(output, target)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        # estimate the accuracy\n",
    "        prediction = output.argmax(-1)\n",
    "        correct = prediction.eq(target.view_as(prediction)).sum()\n",
    "        accuracy = (100.0 * correct / len(target))\n",
    "        accuracy_list.append(accuracy.item())\n",
    "    print('Epoch [{}] loss: {:.5f}; accuracy: {:.2f}%' \\\n",
    "            .format(epoch+1, loss.item(), sum(accuracy_list)/len(accuracy_list)))          "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now test both backbone and linear layer on the test set of CIFAR-10."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test accuracy: 52.33%\n"
     ]
    }
   ],
   "source": [
    "accuracy_list = list()\n",
    "for i, (data, target) in enumerate(test_loader_lineval):\n",
    "    output = backbone_lineval(data).detach()\n",
    "    output = linear_layer(output)\n",
    "    # estimate the accuracy\n",
    "    prediction = output.argmax(-1)\n",
    "    correct = prediction.eq(target.view_as(prediction)).sum()\n",
    "    accuracy = (100.0 * correct / len(target))\n",
    "    accuracy_list.append(accuracy.item())\n",
    "\n",
    "print('Test accuracy: {:.2f}%'.format(sum(accuracy_list)/len(accuracy_list)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the test accuracy is significantly above chance level (in CIFAR-10 chance level is 10%), meaning that during the self-supervised traing it has been possible to build useful representations without accessing the labels."
   ]
  }
 ],
 "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}