report.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DL Project Report\n",
    "Authors: **Giovanni Valer, Emanuele Poiana**\n",
    "\n",
    "## Table of Contents\n",
    "- Introduction\n",
    "- Utils\n",
    "- Data\n",
    "- Entropy Loss\n",
    "    - Weighted Entropy Loss\n",
    "    - Cut Entropy Loss\n",
    "- MEMO\n",
    "    - Standard MEMO\n",
    "    - MEMO with final prediction on marginal distribution\n",
    "    - Batch Normalization\n",
    "- Test Time Adaptation\n",
    "- Experiments\n",
    "- Results and Conclusion\n",
    "- Future Work"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "This project focuses on implementing a Test-Time Adaptation technique for image classification. The goal is to improve the performance of a pre-trained model on out-of-distribution data, without any knowledge of the test-time data distribution. The method we implement is based on **MEMO** (Marginal Entropy Minimization with One test point). The model we use is ResNet-50, pre-trained on **ImageNet**; the dataset we use for testing is **ImageNet-A** (consisting of images that are misclassified by ResNet).\n",
    "\n",
    "The baseline (i.e., ResNet-50 without any adaptation) achieves an accuracy of 0.03%(v1 weights) and 15.3%(v2 weights) on ImageNet-A, so we choose to use v2 weights. We aim to improve this performance by applying MEMO and possibly other techniques.\n",
    "\n",
    "### MEMO, [Zhang et al. (2021)](https://arxiv.org/abs/2110.09506)\n",
    "MEMO consists in applying a set of augmentations to the test image, collecting the output probability of the pre-trained model for each augmented image, and then undertaking a gradient-based optimization to minimize the entropy of the output distributions. The idea is to fine-tune all the model's parameters (on a single test image) to produce consistent predictions across different augmentations.\n",
    "\n",
    "### Our Contributions\n",
    "We further explore other techniques to improve the performance, either modifying MEMO or undertaking different approaches:\n",
    "- Entropy Loss variants (**Weighted Entropy** Loss, **Cut Entropy** Loss)\n",
    "- Final **Prediction on Marginal Distribution** (either standard or weighted average)\n",
    "- **Batch Normalization**\n",
    "- Different **composition of augmentations**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import pandas as pd\n",
    "import torch.nn.functional as F\n",
    "import torchvision.transforms as T\n",
    "import os\n",
    "import boto3\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "from PIL import Image\n",
    "from io import BytesIO\n",
    "from pathlib import Path\n",
    "from torch.utils.data import Dataset\n",
    "from torchvision.models import resnet50, ResNet50_Weights\n",
    "from torchvision.transforms import v2\n",
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ResNet Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def initialize_resnet(weights='v2'):\n",
    "    if weights == 'v1':\n",
    "        resnet = resnet50(weights=ResNet50_Weights.IMAGENET1K_V1).to(device)\n",
    "    elif weights == 'v2':\n",
    "        resnet = resnet50(weights=ResNet50_Weights.IMAGENET1K_V2).to(device)\n",
    "    return resnet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data\n",
    "We need to remap ImageNet-A classes (based on WordNet IDs) to ImageNet classes, in order to use the pre-trained ResNet model.\n",
    "\n",
    "**Note:** if wget fails, download the file from [here](https://github.com/jo-valer/tta-memo/blob/main/label_mappings.tsv)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Comment out if you already downloaded the label mappings\n",
    "!wget https://raw.githubusercontent.com/jo-valer/tta-memo/main/label_mappings.tsv #--no-check-certificate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "label_mapping = pd.read_csv(\"label_mappings.tsv\", sep=\"\\t\")\n",
    "remap_dict = label_mapping.set_index(\"imagenet_a_label\").to_dict()[\"imagenet_label\"]\n",
    "\n",
    "def remap(imagenet_a_label):\n",
    "    \"\"\"\n",
    "    Maps ImageNet-A labels (derived from WordNet IDs) to ImageNet labels (used by ResNet).\n",
    "    \"\"\"\n",
    "    # Map the label\n",
    "    return remap_dict[imagenet_a_label]\n",
    "\n",
    "# We will remove all other labels from the model output\n",
    "labels_in_imagenet_a = [6, 11, 13, 15, 17, 22, 23, 27, 30, 37, 39, 42, 47, 50, 57, 70, 71, 76, 79, 89, 90, 94, 96, 97, 99, 105, 107, 108, 110, 113, 124, 125, 130, 132, 143, 144, 150, 151, 207, 234, 235, 254, 277, 283, 287, 291, 295, 298, 301, 306, 307, 308, 309, 310, 311, 313, 314, 315, 317, 319, 323, 324, 326, 327, 330, 334, 335, 336, 347, 361, 363, 372, 378, 386, 397, 400, 401, 402, 404, 407, 411, 416, 417, 420, 425, 428, 430, 437, 438, 445, 456, 457, 461, 462, 470, 472, 483, 486, 488, 492, 496, 514, 516, 528, 530, 539, 542, 543, 549, 552, 557, 561, 562, 569, 572, 573, 575, 579, 589, 606, 607, 609, 614, 626, 627, 640, 641, 642, 643, 658, 668, 677, 682, 684, 687, 701, 704, 719, 736, 746, 749, 752, 758, 763, 765, 768, 773, 774, 776, 779, 780, 786, 792, 797, 802, 803, 804, 813, 815, 820, 823, 831, 833, 835, 839, 845, 847, 850, 859, 862, 870, 879, 880, 888, 890, 897, 900, 907, 913, 924, 932, 933, 934, 937, 943, 945, 947, 951, 954, 956, 957, 959, 971, 972, 980, 981, 984, 986, 987, 988]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_root=\"imagenet-a\"\n",
    "\n",
    "class ImageNetA(Dataset):\n",
    "    def __init__(self, root, transform=None):\n",
    "        self.s3_bucket = \"deeplearning2024-datasets\"\n",
    "        self.s3_region = \"eu-west-1\"\n",
    "        self.s3_client = boto3.client(\"s3\", region_name=self.s3_region, verify=True)\n",
    "        self.transform = transform\n",
    "\n",
    "        # Get list of objects in the bucket\n",
    "        response = self.s3_client.list_objects_v2(Bucket=self.s3_bucket, Prefix=root)\n",
    "        objects = response.get(\"Contents\", [])\n",
    "        while response.get(\"NextContinuationToken\"):\n",
    "            response = self.s3_client.list_objects_v2(\n",
    "                Bucket=self.s3_bucket,\n",
    "                Prefix=root,\n",
    "                ContinuationToken=response[\"NextContinuationToken\"]\n",
    "            )\n",
    "            objects.extend(response.get(\"Contents\", []))\n",
    "\n",
    "        # Iterate and keep valid files only\n",
    "        self.instances = []\n",
    "        for ds_idx, item in enumerate(objects):\n",
    "            key = item[\"Key\"]\n",
    "            path = Path(key)\n",
    "            \n",
    "            # Check if file is valid\n",
    "            if path.suffix.lower() not in (\".jpg\", \".jpeg\", \".png\", \".ppm\", \".bmp\", \".pgm\", \".tif\", \".tiff\", \".webp\"):\n",
    "                continue\n",
    "\n",
    "            # Get label\n",
    "            label = path.parent.name\n",
    "\n",
    "            # Map the label\n",
    "            label = remap(label)\n",
    "\n",
    "            # Keep track of valid instances\n",
    "            self.instances.append((label, key))\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.instances)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        try:\n",
    "            label, key = self.instances[idx]\n",
    "\n",
    "            label = labels_in_imagenet_a.index(label)\n",
    "\n",
    "            img_bytes = BytesIO()\n",
    "            response = self.s3_client.download_fileobj(Bucket=self.s3_bucket, Key=key, Fileobj=img_bytes)\n",
    "            \n",
    "            # Open image with PIL\n",
    "            img = Image.open(img_bytes).convert(\"RGB\")\n",
    "\n",
    "            # Apply transformations if any\n",
    "            if self.transform is not None:\n",
    "                img = self.transform(img)\n",
    "        except Exception as e:\n",
    "            raise RuntimeError(f\"Error loading image at index {idx}: {str(e)}\")\n",
    "\n",
    "        return img, label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_data(batch_size, img_root, preprocess = False):\n",
    "    # Prepare data transformations for the train loader\n",
    "    transform = None\n",
    "    if preprocess:\n",
    "        transform = T.Compose([\n",
    "            T.Resize((256,256)),\n",
    "            T.CenterCrop((224,224)),\n",
    "            T.ToTensor(),\n",
    "            T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])      # Normalize with ImageNet mean\n",
    "        ])\n",
    "    else:\n",
    "        transform = T.Compose([\n",
    "            T.ToTensor(),\n",
    "            T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "        ])\n",
    "\n",
    "    # Load data\n",
    "    imageneta_dataset = ImageNetA(root=img_root, transform=transform)\n",
    "\n",
    "    # Initialize dataloader\n",
    "    loader = torch.utils.data.DataLoader(imageneta_dataset, batch_size, shuffle=False, num_workers=4)\n",
    "\n",
    "    return loader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Entropy Loss\n",
    "We first implement the Entropy Loss function proposed by [Zhang et al. (2021)](https://arxiv.org/abs/2110.09506). The loss is computed on the outputs of the pre-trained model, after applying the augmentations and is defined as:\n",
    "$$ L_{\\text{E}} = - \\sum_{i=1}^{N} p_i \\log p_i $$\n",
    "where $p_i$ is the probability of class $i$ (averaged over the augmentations) and $N$ is the number of classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EntropyLoss(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super(EntropyLoss, self).__init__()\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Softmax to get probabilities\n",
    "        x = F.softmax(x, dim=1)\n",
    "        # Marginal distribution averaged over augmentations\n",
    "        avg = torch.mean(x, dim=0)\n",
    "        # Entropy loss\n",
    "        return -torch.sum(avg * torch.log(avg))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Weighted Entropy Loss\n",
    "We also implement a variant: the Weighted Entropy Loss. It consists in computing the marginal distribution as a weighted average of the output probabilities over the augmentations.\n",
    "The idea is to give more importance to the augmentations that produce less uncertain predictions, thus the weight associated to each augmentation $a \\in A$ is defined as the inverse of the entropy:\n",
    "$$ w'_a = \\frac{1}{H(p_a)} $$\n",
    "We further normalize the weights over the $A$ augmentations, so that they sum to 1:\n",
    "$$ w_a = \\frac{w'_a}{\\sum_{j=1}^{A} w'_j} $$\n",
    "\n",
    "The marginal distribution is computed as:\n",
    "$$ p_{\\text{m}i} = \\sum_{j=1}^{A} w_j p_{ji} $$\n",
    "where $p_{ji}$ is the probability of class $i$ for the $j$-th augmentation.\n",
    "\n",
    "Finally, the Weighted Entropy Loss is then defined as:\n",
    "$$ L_{\\text{WE}} = - \\sum_{i=1}^{N} p_{\\text{m}i} \\log p_{\\text{m}i} $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class WeightedEntropyLoss(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super(WeightedEntropyLoss, self).__init__()\n",
    "        \n",
    "    def forward(self, x): # A augentations, N classes\n",
    "        # Softmax to get probabilities\n",
    "        probs = F.softmax(x, dim=1) # [A x N]\n",
    "        with torch.no_grad():\n",
    "            # Entropies of the output probabilities\n",
    "            entropies = -torch.sum(probs * torch.log(probs), dim=1) # [A]\n",
    "            # Weights are the inverse of the entropy, normalized\n",
    "            weights = 1 / entropies # [A]\n",
    "            weights /= torch.sum(weights) # [A]\n",
    "        # Weighted probabilities\n",
    "        weighted_probs = weights[:, None] * probs # [A] x [A x N] = [A x N]\n",
    "        # Marginal distribution averaged over augmentations\n",
    "        avg = torch.mean(weighted_probs, dim=0) # [N]\n",
    "        # Entropy loss\n",
    "        return -torch.sum(avg * torch.log(avg))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cut Entropy Loss\n",
    "Another variant we implement is the Cut Entropy Loss. It consists in computing the marginal distribution as the average of the output probabilities over the augmentations, but only considering the top-$k$ augmentations with the lowest entropy. The idea is to avoid handcrafting the set of augmentations to consider, but rather automatically select the most informative ones. The Cut Entropy Loss is defined as:\n",
    "$$ L_{\\text{CE}} = - \\sum_{i=1}^{N} p_{\\text{m}i} \\log p_{\\text{m}i} $$\n",
    "where $p_{\\text{m}i}$ is computed as for the standard Entropy Loss, but only considering the top $k$ augmentations with the lowest entropy:\n",
    "$$ p_{\\text{m}i} = \\sum_{j=1}^{k} w_j p_{ji} $$\n",
    "where:\n",
    "$$ w_j = \\begin{cases} \\frac{1}{k} & \\text{if } j \\leq k \\\\ 0 & \\text{otherwise} \\end{cases} $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CutEntropyLoss(torch.nn.Module):\n",
    "    def __init__(self, cut=0.9):\n",
    "        super(CutEntropyLoss, self).__init__()\n",
    "        self.cut = cut\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Softmax to get probabilities\n",
    "        probs = F.softmax(x, dim=1)\n",
    "        with torch.no_grad():\n",
    "            # Entropy of the output probabilities\n",
    "            entropy = -torch.sum(probs * torch.log(probs), dim=1)\n",
    "            # Sort the entropies\n",
    "            sorted_entropy, _ = torch.sort(entropy, descending=True)\n",
    "            # Assign the weights: 1 for the top cut, 0 for the rest\n",
    "            weights = torch.zeros_like(entropy)\n",
    "            weights[entropy <= sorted_entropy[int(self.cut * len(entropy))]] = 1 # Note: \"top\" here means \"low entropy\"\n",
    "            # Normalize the weights\n",
    "            weights /= torch.sum(weights)\n",
    "        # Weighted marginal distribution\n",
    "        avg = torch.sum(weights[:, None] * probs, dim=0) # [A] x [A x N] = [N]\n",
    "        # Entropy loss\n",
    "        return -torch.sum(avg * torch.log(avg))\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MEMO\n",
    "We implement MEMO as a class wrapping the pre-trained model, the optimizer, and defining the set of augmentations to apply to the test image. We use PyTorch's SGD and perform a single optimization step on the test image, as from preliminary experiments more steps do not seem to improve the performance.\n",
    "\n",
    "We conjointly implement also some variants/alternatives to MEMO, controlled by three parameters:\n",
    "- `memo_train`: whether to train the model to minimize the Entropy Loss on the test image across the augmentations (default: **True**);\n",
    "- `use_augmentations`: instead of making the final prediction on the original data point, make it on the average of the predictions over the augmentations (default: **False**);\n",
    "- `use_weighted_aug`: whether to use the weighted average of the predictions over the augmentations, same as for the Weighted Entropy Loss (default: **False**).\n",
    "\n",
    "The default behavior corresponds to the standard MEMO, i.e., training the model to minimize the Entropy Loss across the augmentations and making the final prediction on the original data point.\n",
    "\n",
    "### Augmentations\n",
    "We experiment with different sets of augmentations, including:\n",
    "- **RandomHorizontalFlip**\n",
    "- **RandomResizedCrop**\n",
    "- **AugMix**\n",
    "\n",
    "### Input Preprocessing\n",
    "The `preprocess` flag specifies the transformation applied in the dataloader for the dataset.\n",
    "- if **True**: Resize and CenterCrop to every image in the set;\n",
    "- if **False**: no transformations are applied (except for the normalization). This requires that the batch size be 1.\n",
    "\n",
    "Such flag is useful to test the model on the original images of the dataset, instead of the resized and cropped ones, so that the augmentations are applied on the original images.\n",
    "\n",
    "### Batch Normalization\n",
    "[Zhang et al. (2021)](https://arxiv.org/abs/2110.09506) suggest to use Batch Normalization on the test image's augmentations, and, to prevent overfitting, to only slighlty update the running statistics of the BatchNorm layers. We implement this technique by setting the `use_adaptive_bn` flag to **True**. This also requires the batch size to be 1, as the running statistics are updated for each image separately. The running statistics are updated as follows:\n",
    "$$ \\mu = \\alpha \\mu_{\\text{train}} + (1 - \\alpha) \\mu_{\\text{test}} $$\n",
    "$$ \\sigma^2 = \\alpha \\sigma^2_{\\text{train}} + (1 - \\alpha) \\sigma^2_{\\text{test}} $$\n",
    "where $\\mu_{\\text{train}}$ and $\\sigma^2_{\\text{train}}$ are the running statistics of the BatchNorm layer on the training set, and $\\mu_{\\text{test}}$ and $\\sigma^2_{\\text{test}}$ are the statistics computed on the test image's augmentations. The parameter $\\alpha$ is set to the best value found in their experiments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def adaptive_bn(self, input):\n",
    "    if self.adapt:\n",
    "        input_mean = torch.zeros(self.running_mean.shape, device=self.running_mean.device)\n",
    "        input_var = torch.ones(self.running_var.shape, device=self.running_var.device)\n",
    "        \n",
    "        # Compute the current mean and variance and put it in \n",
    "        torch.nn.functional.batch_norm(input, input_mean, input_var, None, None, True, 1.0, 0.0)\n",
    "        \n",
    "        # The value 0.9411 is taken directly from the MEMO paper, 16/17 = 0.9411 defined as optimal\n",
    "        adapted_mean = (1- 0.9411) * input_mean + 0.9411 * self.running_mean \n",
    "        adapted_var =  (1- 0.9411) * input_var + 0.9411 * self.running_var\n",
    "        \n",
    "        # Compute the batch_norm with the adapted mean and variance\n",
    "        return torch.nn.functional.batch_norm(input, adapted_mean, adapted_var, self.weight, self.bias, False, 0.0, self.eps)\n",
    "\n",
    "    else:\n",
    "        return torch.nn.functional.batch_norm(input, self.running_mean, self.running_var, self.weight, self.bias, False, 0.1, self.eps)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Implementation\n",
    "All the aforementioned alternatives can be combined with the previously defined Entropy Loss variants, by just passing it as an argument to the MEMO class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Memo(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    Implements the Test-Time Adaptation method of MEMO (Marginal Entropy Minimization with One test point) for image classification.\n",
    "    \"\"\"\n",
    "    def __init__(self, net=resnet50(weights=ResNet50_Weights.IMAGENET1K_V2), optimizer=torch.optim.SGD, lr=0.005, n=64, augmentations='augmix_flip_crop', preprocess=False, device='cuda:0', memo_train:bool=True, use_augmentations:bool=False, use_weighted_aug:bool=False, use_adaptive_bn = False, criterion=EntropyLoss()):\n",
    "        super(Memo, self).__init__()\n",
    "        self.device = device\n",
    "        self.augmentations_name = augmentations # name of augmentations to apply\n",
    "        self.n = n # number of augmentations\n",
    "        self.memo_train = memo_train # whether to train the model at test time or not\n",
    "        self.use_augmentations = use_augmentations # if True, the output is the marginal distribution of the augmented images\n",
    "        self.use_weighted_aug = use_weighted_aug if use_augmentations else False # if True, the output is the weighted marginal distribution of the augmented images (weights for each augmentation are 1/entropy)\n",
    "        self.adapt_bn = use_adaptive_bn\n",
    "        torch.nn.BatchNorm2d.adapt = self.adapt_bn # if True, the batch normalization adapt over input test samples for every BN layer\n",
    "        torch.nn.BatchNorm2d.forward = adaptive_bn \n",
    "        self.net = net.to(device)\n",
    "\n",
    "        self.preprocess = preprocess # if True, the dataloader return the images already preprocessed (cropped, resized, etc.)\n",
    "        self.criterion = criterion\n",
    "\n",
    "        self.optimizer = optimizer(self.net.parameters(), lr=lr)\n",
    "        self.net.eval()\n",
    "\n",
    "        # Save the default parameters of the model\n",
    "        self.default_params = {}\n",
    "        for param in self.net.parameters():\n",
    "            self.default_params[param] = param.clone()\n",
    "\n",
    "        self.cropping = T.Compose([\n",
    "            T.Resize((256,256)),\n",
    "            T.CenterCrop((224,224))\n",
    "        ])\n",
    "\n",
    "        if augmentations == 'augmix':\n",
    "            self.augmentations = v2.AugMix()\n",
    "        elif augmentations == 'augmix_flip':\n",
    "            self.augmentations = T.Compose([v2.AugMix(),\n",
    "                                            T.RandomHorizontalFlip()])\n",
    "        elif augmentations == 'augmix_flip_crop':\n",
    "            self.augmentations = T.Compose([v2.AugMix(),\n",
    "                                            T.RandomHorizontalFlip(),\n",
    "                                            T.RandomResizedCrop(224)])\n",
    "        elif augmentations == 'flip_crop_augmix':\n",
    "            self.augmentations = T.Compose([T.RandomHorizontalFlip(),\n",
    "                                            T.RandomResizedCrop(224),\n",
    "                                            v2.AugMix()])\n",
    "        elif augmentations == 'crop':\n",
    "            self.augmentations = T.Compose([T.RandomResizedCrop(224)])\n",
    "        elif augmentations == 'flip_crop':\n",
    "            self.augmentations = T.Compose([T.RandomHorizontalFlip(),\n",
    "                                            T.RandomResizedCrop(224)])\n",
    "        else: # Throw error\n",
    "            raise ValueError(\"Invalid augmentations. Choose from 'augmix', 'augmix_flip', 'augmix_flip_crop', 'flip_crop_augmix', 'crop', 'flip_crop'\")\n",
    "    \n",
    "    def filter_classes(self, outputs):\n",
    "        return outputs[:, labels_in_imagenet_a]\n",
    "\n",
    "    def apply_augmentations(self, img_original, augmentations):\n",
    "        \"\"\"\n",
    "        Applies the augmentations to the input image.\n",
    "        \"\"\"\n",
    "        with torch.no_grad():\n",
    "            return augmentations(img_original)\n",
    "    \n",
    "    def get_augment_loader(self, x):\n",
    "        \"\"\"\n",
    "        Returns the dataloader with the input image x and its augmentations.\n",
    "        If the input image is already cropped, or the augmentations do not apply cropping, the image does not require cropping.\n",
    "        \"\"\"\n",
    "        with torch.no_grad():\n",
    "            if self.preprocess or (self.augmentations_name == 'augmix_flip' or self.augmentations_name == 'augmix'):\n",
    "                aug_images = [x]\n",
    "            else:\n",
    "                aug_images = [self.cropping(x)] # This ensures that the original image and the augmented images are of the same size\n",
    "        \n",
    "            for _ in range(self.n):\n",
    "                aug_img = self.apply_augmentations(x, augmentations=self.augmentations)#.to(device)\n",
    "                aug_images.append(aug_img)\n",
    "        \n",
    "            return torch.utils.data.DataLoader(aug_images, batch_size = self.n + 1, shuffle=False, num_workers=0)\n",
    "\n",
    "    def augment(self, x):\n",
    "        \"\"\"\n",
    "        Returns the image x and its augmentations.\n",
    "        \"\"\"\n",
    "        aug_loader = self.get_augment_loader(x)\n",
    "        \n",
    "        # Get a single batch of augmented images\n",
    "        aug_inputs = next(iter(aug_loader)).to(self.device)\n",
    "\n",
    "        return aug_inputs\n",
    "    \n",
    "    def get_averaged_distribution(self, x):\n",
    "        \"\"\"\n",
    "        Returns the averaged marginal distribution of the network over the input image and its augmentations.\n",
    "        If use_weighted_aug is True, the output is the weighted marginal distribution of the augmented images (weights for each augmentation are 1/entropy).\n",
    "        \"\"\"\n",
    "        predictions = None\n",
    "        for image in x:\n",
    "            inputs = self.augment(image)\n",
    "            outputs = self.filter_classes(self.net(inputs))\n",
    "\n",
    "            if self.use_weighted_aug:\n",
    "                # Entropy of the output probabilities (softmax of the network output)\n",
    "                entropies = -torch.sum(F.softmax(outputs, dim=1) * torch.log(F.softmax(outputs, dim=1)), dim=1)\n",
    "                weights = 1 / entropies\n",
    "                weights /= torch.sum(weights)\n",
    "                weighted_outputs = weights[:, None] * outputs\n",
    "                mean_distribution = torch.mean(weighted_outputs, dim=0).unsqueeze(dim=0)\n",
    "            else:\n",
    "                mean_distribution = torch.mean(outputs, dim=0).unsqueeze(dim=0)\n",
    "            if predictions is None:\n",
    "                predictions = mean_distribution\n",
    "            else:\n",
    "                predictions = torch.cat((predictions, mean_distribution), 0)\n",
    "\n",
    "        return predictions\n",
    "\n",
    "    def test_time_training(self, x):\n",
    "        \"\"\"\n",
    "        Trains the model at test time.\n",
    "        \"\"\"\n",
    "        predictions = None\n",
    "        for image in x:\n",
    "            torch.cuda.empty_cache()\n",
    "            self.optimizer.zero_grad()\n",
    "\n",
    "            inputs = self.augment(image)\n",
    "            outputs = self.filter_classes(self.net(inputs))\n",
    "\n",
    "            del inputs\n",
    "            \n",
    "            # Compute the entropy loss\n",
    "            entropy_loss = self.criterion(outputs)\n",
    "            \n",
    "            # Backpropagate the loss\n",
    "            entropy_loss.backward()\n",
    "            \n",
    "            # Update the weights\n",
    "            self.optimizer.step()\n",
    "            \n",
    "            if not self.preprocess and not (self.augmentations_name == 'augmix_flip' or self.augmentations_name == 'augmix'):\n",
    "                image = self.cropping(image)\n",
    "            \n",
    "            if not self.use_augmentations:\n",
    "                prediction = self.filter_classes(self.net(image.unsqueeze(dim=0)))\n",
    "            else:\n",
    "                prediction = self.get_averaged_distribution(image.unsqueeze(dim=0))\n",
    "\n",
    "            if predictions is None:\n",
    "                predictions = prediction\n",
    "            else:\n",
    "                predictions = torch.cat((predictions, prediction), 0)\n",
    "\n",
    "            # Restore the default parameters\n",
    "            for param in self.net.parameters():\n",
    "                param.data = self.default_params[param].clone()\n",
    "        \n",
    "        return predictions\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Forward pass of the model, eventually training the model at test time.\n",
    "        \"\"\"\n",
    "        if self.memo_train:\n",
    "            output = self.test_time_training(x)\n",
    "        else:\n",
    "            if self.use_augmentations:\n",
    "                output = self.get_averaged_distribution(x)\n",
    "            else:\n",
    "                output = self.filter_classes(self.net(x))\n",
    "\n",
    "        return output\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test Time Adaptation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(net, data_loader, device=\"cuda\"):\n",
    "    samples = 0.0\n",
    "    cumulative_accuracy = 0.0\n",
    "\n",
    "    # Set the network to evaluation mode\n",
    "    net.eval()\n",
    "\n",
    "    # Iterate over the test set\n",
    "    for _, (inputs, targets) in enumerate(tqdm(data_loader)):\n",
    "\n",
    "        # Load data into GPU\n",
    "        inputs = inputs.to(device)\n",
    "        targets = targets.to(device)\n",
    "\n",
    "        # Forward pass\n",
    "        outputs = net(inputs)\n",
    "\n",
    "        # Disable gradient computation (we are only testing, we do not want our model to be modified in this step!)\n",
    "        with torch.no_grad():\n",
    "\n",
    "            # Fetch prediction\n",
    "            samples += inputs.shape[0]\n",
    "            _, predicted = outputs.max(1)\n",
    "\n",
    "            # Compute accuracy\n",
    "            cumulative_accuracy += predicted.eq(targets).sum().item()\n",
    "            del outputs, inputs, targets, predicted\n",
    "\n",
    "    return cumulative_accuracy / samples * 100"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Experiments\n",
    "We hereby present a selection of significant experiments.\n",
    "\n",
    "### MEMO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = Memo(\n",
    "    net=initialize_resnet(),\n",
    "    device=device\n",
    "    )\n",
    "\n",
    "loader = get_data(1, img_root)\n",
    "test_accuracy = test(net, loader, device=device)\n",
    "print(f\"\\tTop-1 accuracy {test_accuracy:.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### MEMO + Cut Entropy Loss + Batch Normalization + Final Prediction on Marginal Distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = Memo(\n",
    "    net=initialize_resnet(),\n",
    "    augmentations=\"crop\",\n",
    "    use_augmentations=True,\n",
    "    criterion=CutEntropyLoss(cut=0.9),\n",
    "    use_adaptive_bn=True,\n",
    "    device=device\n",
    "    )\n",
    "\n",
    "loader = get_data(1, img_root)\n",
    "test_accuracy = test(net, loader, device=device)\n",
    "print(f\"\\tTop-1 accuracy {test_accuracy:.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### MEMO + CEL + BN + Final Prediction on Weighted Marginal Distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = Memo(\n",
    "    net=initialize_resnet(),\n",
    "    augmentations=\"crop\",\n",
    "    use_augmentations=True,\n",
    "    use_weighted_aug=True,\n",
    "    criterion=CutEntropyLoss(cut=0.9),\n",
    "    use_adaptive_bn=True,\n",
    "    device=device\n",
    "    )\n",
    "\n",
    "loader = get_data(1, img_root)\n",
    "test_accuracy = test(net, loader, device=device)\n",
    "print(f\"\\tTop-1 accuracy {test_accuracy:.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Average Output Probabilities of Augmentations\n",
    "We also experiment by avoiding any training steps: we use the averaged output probabilities (over the augmentations) as the final prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = Memo(\n",
    "    net=initialize_resnet(),\n",
    "    augmentations=\"flip_crop\",\n",
    "    memo_train=False,\n",
    "    use_augmentations=True,\n",
    "    use_adaptive_bn=True,\n",
    "    device=device\n",
    "    )\n",
    "\n",
    "loader = get_data(1, img_root)\n",
    "test_accuracy = test(net, loader, device=device)\n",
    "print(f\"\\tTop-1 accuracy {test_accuracy:.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Weighted Average Output Probabilities of Augmentations\n",
    "Similarly, we experiment with the weighted average of the output probabilities, using the weights computed based on the entropy of the output probabilities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = Memo(\n",
    "    net=initialize_resnet(),\n",
    "    augmentations=\"crop\",\n",
    "    memo_train=False,\n",
    "    use_augmentations=True,\n",
    "    use_weighted_aug=True,\n",
    "    use_adaptive_bn=True,\n",
    "    device=device\n",
    "    )\n",
    "\n",
    "loader = get_data(1, img_root)\n",
    "test_accuracy = test(net, loader, device=device)\n",
    "print(f\"\\tTop-1 accuracy {test_accuracy:.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Results\n",
    "\n",
    "The baseline model achieves a top-1 accuracy of **15.07%**, when restricting the output labels to those in Imagenet-A.\n",
    "We find that MEMO improves the performance of the pre-trained model on ImageNet-A, and that combining it other techniques can further improve the performance.\n",
    "\n",
    "The best results of MEMO are obtained with with 64 augmentations and a learning rate of 0.005. When using the augmentations proposed by [Zhang et al. (2021)](https://arxiv.org/abs/2110.09506) we find a top-1 accuracy of **16.17%**; while with RandomResizedCrop we achieve **19.87%**.\n",
    "In the following, we present the results in an incremental fashion, showing the impact of each technique on the performance.\n",
    "\n",
    "### Entropy Loss\n",
    "From preliminary experiments we notice that the Weighted Entropy Loss does not improve the performance, compared to the standard Entropy Loss. We therefore focus on the **Cut Entropy Loss**, which instead has a **beneficial impact on the accuracy** of the model: the top-$k$ augmentations (those with lower entropy) are much more informative than the handcrafted set of augmentations. We find the best results with the top-10% and top-20% (i.e. cut=0.9 and cut=0.8, respectively), out of 64 augmentations. This further improves the accuracy to **23.20%**.\n",
    "\n",
    "### Augmentations\n",
    "We tried different sets of augmentations, and we observed that random cropping part of the original (full size) image produces the best results, together with random horizontal flipping.\n",
    "This is probably due to the characteristics of the ImageNet-A dataset, where the true labels of the images are usually located in small, sparse areas of the test sample.\n",
    "Random cropping prompts the model to avoid the “bigger” objects in the image and improves the probability of detecting smaller but more relevant features.\n",
    "Speaking of top-1 accuracy, the best results to this point are indeed obtained with RandomResizedCrop+RandomHorizontalFlip: **23.20%**, while adding AugMix leads to **19.86%**.\n",
    "\n",
    "### Batch Normalization Adaptation\n",
    "Batch Normalization is implemented in ResNet50; but, as suggested by [Schneider et al. (2020)](https://arxiv.org/abs/2006.16971), just a single test sample adaptation is enough to generally show an improvement in the model performance.\n",
    "We achieve better results with BN, in particular it further improves the accuracy to **27.68%**, using a learning rate of 0.01.\n",
    "\n",
    "### Final prediction on marginal distribution\n",
    "This approach can be used either alone or in combination with MEMO. Both when added on top of previous techniques and when used alone, it achieves similar results: **27.68%**. It is worth noting, however, that is requires less computational resources, and thus proves to be the most promising technique. We also expected to achieve even better results, but due to time and computational constraints we have not been able to properly search the best hyperparameters' values (as there is a considerable number of possible combinations).\\\n",
    "The alternative of using a weighted average, instead of the standard average of the output probabilities over the augmentations, does not seem to be particularly effective, as it does not improve the accuracy.\n",
    "\n",
    "\n",
    "## Conclusion\n",
    "We showed how MEMO can be a viable TTA technique, but other approaches are more effective and also faster, at least in our experimental setup. In particular, the best results are obtained by using the MEMO with RandomResizedCrop, Cut Entropy Loss, and Batch Normalization. The final prediction on the marginal distribution also seems to be a promising technique, reaching the same best results in our experiments.\n",
    "\n",
    "### Future Work\n",
    "Due to time and computational constraints, we could not undertake more experiments; however, we foresee another method for test-time adaptation, which can be either combined with MEMO or used as an alternative: **APoZ-based Adaptation**. The idea is to prune the model based on the Average Percentage of Zeros (APoZ) of the activations (before the final fully connected layer) over the augmentations of the test image. The intuition is that the nodes that are more active across the augmentations are more informative for the target class. By pruning the model, we are retaining only the most informative features, and therefore adapting the model to the test-time data point.\n"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}