examples/notebooks/EfficientNet_Cifar100_finetuning.ipynb

{
 "nbformat": 4,
 "nbformat_minor": 0,
 "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.7.4"
  },
  "colab": {
   "name": "EfficientNet_Cifar100_finetuning.ipynb",
   "provenance": []
  },
  "accelerator": "GPU"
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JZE22y3LsXpd"
   },
   "source": [
    "# Finetuning of ImageNet pretrained EfficientNet-B0 on CIFAR-100\n",
    "\n",
    "In 2019, new ConvNets architectures have been proposed in [\"EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks\"](https://arxiv.org/pdf/1905.11946.pdf) paper. According to the paper, model's compound scaling starting from a 'good' baseline provides an network that achieves  state-of-the-art on  ImageNet,  while  being 8.4x  smaller and 6.1x faster on inference than the best existing ConvNet.\n",
    "\n",
    "![efficientnets](https://github.com/abdulelahsm/ignite/blob/update-tutorials/examples/notebooks/assets/efficientnets.png?raw=1)\n",
    "\n",
    "Following the paper, EfficientNet-B0 model pretrained on ImageNet and finetuned on CIFAR100 dataset gives 88% test accuracy. Let's reproduce this result with Ignite. [Official implementation](https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet) of EfficientNet uses Tensorflow, \n",
    "for our case we will borrow the code from [katsura-jp/efficientnet-pytorch](https://github.com/katsura-jp/efficientnet-pytorch), \n",
    "[rwightman/pytorch-image-models](https://github.com/rwightman/pytorch-image-models) and [lukemelas/EfficientNet-PyTorch](https://github.com/lukemelas/EfficientNet-PyTorch/) repositories (kudos to authors!). We will download pretrained weights from [lukemelas/EfficientNet-PyTorch](https://github.com/lukemelas/EfficientNet-PyTorch/) repository.\n",
    "\n",
    "## Network architecture review\n",
    "The architecture of EfficientNet-B0 is the following:\n",
    "```\n",
    "1 - Stem    - Conv3x3|BN|Swish\n",
    "\n",
    "2 - Blocks  - MBConv1, k3x3 \n",
    "            - MBConv6, k3x3 repeated 2 times\n",
    "            - MBConv6, k5x5 repeated 2 times\n",
    "            - MBConv6, k3x3 repeated 3 times\n",
    "            - MBConv6, k5x5 repeated 3 times\n",
    "            - MBConv6, k5x5 repeated 4 times\n",
    "            - MBConv6, k3x3\n",
    "                            totally 16 blocks\n",
    "\n",
    "3 - Head    - Conv1x1|BN|Swish \n",
    "            - Pooling\n",
    "            - Dropout\n",
    "            - FC\n",
    "```\n",
    "\n",
    "where \n",
    "```\n",
    "Swish(x) = x * sigmoid(x)\n",
    "```\n",
    "and `MBConvX` stands for mobile inverted bottleneck convolution, X - denotes expansion ratio:\n",
    "``` \n",
    "MBConv1 : \n",
    "  -> DepthwiseConv|BN|Swish -> SqueezeExcitation -> Conv|BN\n",
    "\n",
    "MBConv6 : \n",
    "  -> Conv|BN|Swish -> DepthwiseConv|BN|Swish -> SqueezeExcitation -> Conv|BN\n",
    "\n",
    "MBConv6+IdentitySkip : \n",
    "  -.-> Conv|BN|Swish -> DepthwiseConv|BN|Swish -> SqueezeExcitation -> Conv|BN-(+)->\n",
    "   \\___________________________________________________________________________/\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hP_tseP1sXpl"
   },
   "source": [
    "## Installations\n",
    "\n",
    "1) Torchvision\n",
    "\n",
    "Please install torchvision in order to get CIFAR100 dataset: \n",
    "```\n",
    "conda install -y torchvision -c pytorch\n",
    "```\n",
    "\n",
    "2) Let's install Nvidia/Apex package:\n",
    "\n",
    "We will train with automatic mixed precision using [nvidia/apex](https://github.com/NVIDIA/apex) pacakge"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "id": "br990PJfgCHz"
   },
   "source": [
    "# Install Apex:\n",
    "# If torch cuda version and nvcc version match:\n",
    "!pip install --upgrade --no-cache-dir --global-option=\"--cpp_ext\" --global-option=\"--cuda_ext\" git+https://github.com/NVIDIA/apex/\n",
    "# if above command is failing, please install apex without c++/cuda extensions:\n",
    "# !pip install --upgrade --no-cache-dir git+https://github.com/NVIDIA/apex/"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "id": "naLdioPTgCH2"
   },
   "source": [
    "3) Install tensorboardX and `pytorch-ignite`"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "watnTwx-sXpm"
   },
   "source": [
    "!pip install pytorch-ignite tensorboardX"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "FudN7dJBsXpp"
   },
   "source": [
    "import random\n",
    "import torch\n",
    "import ignite\n",
    "\n",
    "seed = 17\n",
    "random.seed(seed)\n",
    "_ = torch.manual_seed(seed)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "EnytGeC5sXpq"
   },
   "source": [
    "torch.__version__, ignite.__version__"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IZEFHfBwsXpr"
   },
   "source": [
    "## Model\n",
    "\n",
    "\n",
    "Let's define some helpful modules:\n",
    "- Flatten \n",
    "- Swish \n",
    "\n",
    "The reason why Swish is not implemented in `torch.nn` can be found [here](https://github.com/pytorch/pytorch/pull/3182).\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "rREuFNq1sXps"
   },
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "\n",
    "class Swish(nn.Module):\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return x * torch.sigmoid(x)\n",
    "\n",
    "\n",
    "class Flatten(nn.Module):\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return x.reshape(x.shape[0], -1)\n"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tlOJJKnVsXpt"
   },
   "source": [
    "Let's visualize Swish transform vs ReLU:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "SiQ5NmqasXpu"
   },
   "source": [
    "import matplotlib.pylab as plt\n",
    "%matplotlib inline\n",
    "\n",
    "d = torch.linspace(-10.0, 10.0)\n",
    "s = Swish()\n",
    "res = s(d)\n",
    "res2 = torch.relu(d)\n",
    "\n",
    "plt.title(\"Swish transformation\")\n",
    "plt.plot(d.numpy(), res.numpy(), label='Swish')\n",
    "plt.plot(d.numpy(), res2.numpy(), label='ReLU')\n",
    "plt.legend()"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "i6Sp-82PsXpv"
   },
   "source": [
    "Now let's define `SqueezeExcitation` module"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "rmvJRZQosXpw"
   },
   "source": [
    "class SqueezeExcitation(nn.Module):\n",
    "    \n",
    "    def __init__(self, inplanes, se_planes):\n",
    "        super(SqueezeExcitation, self).__init__()\n",
    "        self.reduce_expand = nn.Sequential(\n",
    "            nn.Conv2d(inplanes, se_planes, \n",
    "                      kernel_size=1, stride=1, padding=0, bias=True),\n",
    "            Swish(),\n",
    "            nn.Conv2d(se_planes, inplanes, \n",
    "                      kernel_size=1, stride=1, padding=0, bias=True),\n",
    "            nn.Sigmoid()\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x_se = torch.mean(x, dim=(-2, -1), keepdim=True)\n",
    "        x_se = self.reduce_expand(x_se)\n",
    "        return x_se * x\n"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-CAxubs8sXpw"
   },
   "source": [
    "Next, we can define `MBConv`.\n",
    "\n",
    "**Note on implementation**: in Tensorflow (and PyTorch ports) convolutions use `SAME` padding option which in PyTorch requires\n",
    "a specific padding computation and additional operation to apply. We will use built-in padding argument of the convolution."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "6-AL5AYysXpw"
   },
   "source": [
    "from torch.nn import functional as F\n",
    "\n",
    "\n",
    "class MBConv(nn.Module):\n",
    "\n",
    "    def __init__(self, inplanes, planes, kernel_size, stride, \n",
    "                 expand_rate=1.0, se_rate=0.25, \n",
    "                 drop_connect_rate=0.2):\n",
    "        super(MBConv, self).__init__()\n",
    "\n",
    "        expand_planes = int(inplanes * expand_rate)\n",
    "        se_planes = max(1, int(inplanes * se_rate))\n",
    "\n",
    "        self.expansion_conv = None        \n",
    "        if expand_rate > 1.0:\n",
    "            self.expansion_conv = nn.Sequential(\n",
    "                nn.Conv2d(inplanes, expand_planes, \n",
    "                          kernel_size=1, stride=1, padding=0, bias=False),\n",
    "                nn.BatchNorm2d(expand_planes, momentum=0.01, eps=1e-3),\n",
    "                Swish()\n",
    "            )\n",
    "            inplanes = expand_planes\n",
    "\n",
    "        self.depthwise_conv = nn.Sequential(\n",
    "            nn.Conv2d(inplanes, expand_planes,\n",
    "                      kernel_size=kernel_size, stride=stride, \n",
    "                      padding=kernel_size // 2, groups=expand_planes,\n",
    "                      bias=False),\n",
    "            nn.BatchNorm2d(expand_planes, momentum=0.01, eps=1e-3),\n",
    "            Swish()\n",
    "        )\n",
    "\n",
    "        self.squeeze_excitation = SqueezeExcitation(expand_planes, se_planes)\n",
    "        \n",
    "        self.project_conv = nn.Sequential(\n",
    "            nn.Conv2d(expand_planes, planes, \n",
    "                      kernel_size=1, stride=1, padding=0, bias=False),\n",
    "            nn.BatchNorm2d(planes, momentum=0.01, eps=1e-3),\n",
    "        )\n",
    "\n",
    "        self.with_skip = stride == 1\n",
    "        self.drop_connect_rate = drop_connect_rate\n",
    "    \n",
    "    def _drop_connect(self, x):        \n",
    "        keep_prob = 1.0 - self.drop_connect_rate\n",
    "        drop_mask = torch.rand(x.shape[0], 1, 1, 1) + keep_prob\n",
    "        drop_mask = drop_mask.type_as(x)\n",
    "        drop_mask.floor_()\n",
    "        return drop_mask * x / keep_prob\n",
    "        \n",
    "    def forward(self, x):\n",
    "        z = x\n",
    "        if self.expansion_conv is not None:\n",
    "            x = self.expansion_conv(x)\n",
    "\n",
    "        x = self.depthwise_conv(x)\n",
    "        x = self.squeeze_excitation(x)\n",
    "        x = self.project_conv(x)\n",
    "        \n",
    "        # Add identity skip\n",
    "        if x.shape == z.shape and self.with_skip:            \n",
    "            if self.training and self.drop_connect_rate is not None:\n",
    "                x = self._drop_connect(x)\n",
    "            x += z\n",
    "        return x"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OlPTQlFRsXpx"
   },
   "source": [
    "And finally, we can implement generic `EfficientNet`:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "CV2NfxZIsXpx"
   },
   "source": [
    "from collections import OrderedDict\n",
    "import math\n",
    "\n",
    "\n",
    "def init_weights(module):    \n",
    "    if isinstance(module, nn.Conv2d):    \n",
    "        nn.init.kaiming_normal_(module.weight, a=0, mode='fan_out')\n",
    "    elif isinstance(module, nn.Linear):\n",
    "        init_range = 1.0 / math.sqrt(module.weight.shape[1])\n",
    "        nn.init.uniform_(module.weight, a=-init_range, b=init_range)\n",
    "        \n",
    "        \n",
    "class EfficientNet(nn.Module):\n",
    "        \n",
    "    def _setup_repeats(self, num_repeats):\n",
    "        return int(math.ceil(self.depth_coefficient * num_repeats))\n",
    "    \n",
    "    def _setup_channels(self, num_channels):\n",
    "        num_channels *= self.width_coefficient\n",
    "        new_num_channels = math.floor(num_channels / self.divisor + 0.5) * self.divisor\n",
    "        new_num_channels = max(self.divisor, new_num_channels)\n",
    "        if new_num_channels < 0.9 * num_channels:\n",
    "            new_num_channels += self.divisor\n",
    "        return new_num_channels\n",
    "\n",
    "    def __init__(self, num_classes=100, \n",
    "                 width_coefficient=1.0,\n",
    "                 depth_coefficient=1.0,\n",
    "                 se_rate=0.25,\n",
    "                 dropout_rate=0.2,\n",
    "                 drop_connect_rate=0.2):\n",
    "        super(EfficientNet, self).__init__()\n",
    "        \n",
    "        self.width_coefficient = width_coefficient\n",
    "        self.depth_coefficient = depth_coefficient\n",
    "        self.divisor = 8\n",
    "                \n",
    "        list_channels = [32, 16, 24, 40, 80, 112, 192, 320, 1280]\n",
    "        list_channels = [self._setup_channels(c) for c in list_channels]\n",
    "                \n",
    "        list_num_repeats = [1, 2, 2, 3, 3, 4, 1]\n",
    "        list_num_repeats = [self._setup_repeats(r) for r in list_num_repeats]        \n",
    "        \n",
    "        expand_rates = [1, 6, 6, 6, 6, 6, 6]\n",
    "        strides = [1, 2, 2, 2, 1, 2, 1]\n",
    "        kernel_sizes = [3, 3, 5, 3, 5, 5, 3]\n",
    "\n",
    "        # Define stem:\n",
    "        self.stem = nn.Sequential(\n",
    "            nn.Conv2d(3, list_channels[0], kernel_size=3, stride=2, padding=1, bias=False),\n",
    "            nn.BatchNorm2d(list_channels[0], momentum=0.01, eps=1e-3),\n",
    "            Swish()\n",
    "        )\n",
    "        \n",
    "        # Define MBConv blocks\n",
    "        blocks = []\n",
    "        counter = 0\n",
    "        num_blocks = sum(list_num_repeats)\n",
    "        for idx in range(7):\n",
    "            \n",
    "            num_channels = list_channels[idx]\n",
    "            next_num_channels = list_channels[idx + 1]\n",
    "            num_repeats = list_num_repeats[idx]\n",
    "            expand_rate = expand_rates[idx]\n",
    "            kernel_size = kernel_sizes[idx]\n",
    "            stride = strides[idx]\n",
    "            drop_rate = drop_connect_rate * counter / num_blocks\n",
    "            \n",
    "            name = \"MBConv{}_{}\".format(expand_rate, counter)\n",
    "            blocks.append((\n",
    "                name,\n",
    "                MBConv(num_channels, next_num_channels, \n",
    "                       kernel_size=kernel_size, stride=stride, expand_rate=expand_rate, \n",
    "                       se_rate=se_rate, drop_connect_rate=drop_rate)\n",
    "            ))\n",
    "            counter += 1\n",
    "            for i in range(1, num_repeats):                \n",
    "                name = \"MBConv{}_{}\".format(expand_rate, counter)\n",
    "                drop_rate = drop_connect_rate * counter / num_blocks                \n",
    "                blocks.append((\n",
    "                    name,\n",
    "                    MBConv(next_num_channels, next_num_channels, \n",
    "                           kernel_size=kernel_size, stride=1, expand_rate=expand_rate, \n",
    "                           se_rate=se_rate, drop_connect_rate=drop_rate)                                    \n",
    "                ))\n",
    "                counter += 1\n",
    "        \n",
    "        self.blocks = nn.Sequential(OrderedDict(blocks))\n",
    "        \n",
    "        # Define head\n",
    "        self.head = nn.Sequential(\n",
    "            nn.Conv2d(list_channels[-2], list_channels[-1], \n",
    "                      kernel_size=1, bias=False),\n",
    "            nn.BatchNorm2d(list_channels[-1], momentum=0.01, eps=1e-3),\n",
    "            Swish(),\n",
    "            nn.AdaptiveAvgPool2d(1),\n",
    "            Flatten(),\n",
    "            nn.Dropout(p=dropout_rate),\n",
    "            nn.Linear(list_channels[-1], num_classes)\n",
    "        )\n",
    "\n",
    "        self.apply(init_weights)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        f = self.stem(x)\n",
    "        f = self.blocks(f)\n",
    "        y = self.head(f)\n",
    "        return y"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9tEhbJWnsXpy"
   },
   "source": [
    "All EfficientNet models can be defined using the following parametrization:\n",
    "```\n",
    "# (width_coefficient, depth_coefficient, resolution, dropout_rate)\n",
    "'efficientnet-b0': (1.0, 1.0, 224, 0.2),\n",
    "'efficientnet-b1': (1.0, 1.1, 240, 0.2),\n",
    "'efficientnet-b2': (1.1, 1.2, 260, 0.3),\n",
    "'efficientnet-b3': (1.2, 1.4, 300, 0.3),\n",
    "'efficientnet-b4': (1.4, 1.8, 380, 0.4),\n",
    "'efficientnet-b5': (1.6, 2.2, 456, 0.4),\n",
    "'efficientnet-b6': (1.8, 2.6, 528, 0.5),\n",
    "'efficientnet-b7': (2.0, 3.1, 600, 0.5),\n",
    "```    \n",
    "Let's define and train the third one: `EfficientNet-B0`"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "YCuCoTHBsXpy"
   },
   "source": [
    "model = EfficientNet(num_classes=1000, \n",
    "                     width_coefficient=1.0, depth_coefficient=1.0, \n",
    "                     dropout_rate=0.2)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HjUCNwu3sXpz"
   },
   "source": [
    "Number of parameters:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "WfBynpOEsXpz"
   },
   "source": [
    "def print_num_params(model, display_all_modules=False):\n",
    "    total_num_params = 0\n",
    "    for n, p in model.named_parameters():\n",
    "        num_params = 1\n",
    "        for s in p.shape:\n",
    "            num_params *= s\n",
    "        if display_all_modules: print(\"{}: {}\".format(n, num_params))\n",
    "        total_num_params += num_params\n",
    "    print(\"-\" * 50)\n",
    "    print(\"Total number of parameters: {:.2e}\".format(total_num_params))\n",
    "    \n",
    "\n",
    "print_num_params(model)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DxiAjmuMsXp0"
   },
   "source": [
    "Let's compare the number of parameters with some of ResNets:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "rAUg5NwnsXp0"
   },
   "source": [
    "from torchvision.models.resnet import resnet18, resnet34, resnet50"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "9ZMMOAjpsXp0"
   },
   "source": [
    "print_num_params(resnet18(pretrained=False, num_classes=100))\n",
    "print_num_params(resnet34(pretrained=False, num_classes=100))\n",
    "print_num_params(resnet50(pretrained=False, num_classes=100))"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "0CcItj-CsXp0"
   },
   "source": [],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0NhRnYnAsXp1"
   },
   "source": [
    "### Model's graph with Tensorboard\n",
    "\n",
    "We can optionally inspect model's graph with the code below. For that we need to install\n",
    "`tensorboardX` package.\n",
    "Otherwise go directly to the next section."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "F2y5WvBisXp1"
   },
   "source": [
    "from tensorboardX.pytorch_graph import graph\n",
    "\n",
    "import random\n",
    "from IPython.display import clear_output, Image, display, HTML\n",
    "\n",
    "\n",
    "def show_graph(graph_def):\n",
    "    \"\"\"Visualize TensorFlow graph.\"\"\"\n",
    "    if hasattr(graph_def, 'as_graph_def'):\n",
    "        graph_def = graph_def.as_graph_def()\n",
    "    strip_def = graph_def\n",
    "    code = \"\"\"\n",
    "        <script src=\"//cdnjs.cloudflare.com/ajax/libs/polymer/0.3.3/platform.js\"></script>\n",
    "        <script>\n",
    "          function load() {{\n",
    "            document.getElementById(\"{id}\").pbtxt = {data};\n",
    "          }}\n",
    "        </script>\n",
    "        <link rel=\"import\" href=\"https://tensorboard.appspot.com/tf-graph-basic.build.html\" onload=load()>\n",
    "        <div style=\"height:600px\">\n",
    "          <tf-graph-basic id=\"{id}\"></tf-graph-basic>\n",
    "        </div>\n",
    "    \"\"\".format(data=repr(str(strip_def)), id='graph'+str(random.randint(0, 1000)))\n",
    "\n",
    "    iframe = \"\"\"\n",
    "        <iframe seamless style=\"width:1200px;height:620px;border:0\" srcdoc=\"{}\"></iframe>\n",
    "    \"\"\".format(code.replace('\"', '&quot;'))\n",
    "    display(HTML(iframe))"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "scrolled": true,
    "id": "_oOEmRUJsXp1"
   },
   "source": [
    "x = torch.rand(4, 3, 224, 224)\n",
    "\n",
    "# Error : module 'torch.onnx' has no attribute 'set_training'\n",
    "# uncomment when  it will be fixed \n",
    "\n",
    "# graph_def = graph(model, x, operator_export_type='RAW')"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "vW1y3eTMsXp1"
   },
   "source": [
    "# Display in Firefox may not work properly. Use Chrome.\n",
    "# show_graph(graph_def[0])"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "irrGiMKWsXp4"
   },
   "source": [
    "### Load pretrained weights\n",
    "\n",
    "Let's load pretrained weights and check the model on a single image."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "gd3CHmnhsXp4"
   },
   "source": [
    "!mkdir /tmp/efficientnet_weights\n",
    "!wget http://storage.googleapis.com/public-models/efficientnet-b0-08094119.pth -O/tmp/efficientnet_weights/efficientnet-b0-08094119.pth"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "0f1BnJ4TsXp4"
   },
   "source": [
    "from collections import OrderedDict\n",
    "\n",
    "model_state = torch.load(\"/tmp/efficientnet_weights/efficientnet-b0-08094119.pth\")\n",
    "\n",
    "# A basic remapping is required\n",
    "mapping = {\n",
    "    k: v for k, v in zip(model_state.keys(), model.state_dict().keys())\n",
    "}\n",
    "mapped_model_state = OrderedDict([\n",
    "    (mapping[k], v) for k, v in model_state.items()\n",
    "])\n",
    "\n",
    "model.load_state_dict(mapped_model_state, strict=False)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "-9mxclp2sXp5"
   },
   "source": [
    "!wget https://raw.githubusercontent.com/lukemelas/EfficientNet-PyTorch/master/examples/simple/img.jpg -O/tmp/giant_panda.jpg\n",
    "!wget https://raw.githubusercontent.com/lukemelas/EfficientNet-PyTorch/master/examples/simple/labels_map.txt -O/tmp/labels_map.txt"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "3QUPjDRBsXp5"
   },
   "source": [
    "import json\n",
    "\n",
    "with open(\"/tmp/labels_map.txt\", \"r\") as h:\n",
    "    labels = json.load(h)\n",
    "\n",
    "from PIL import Image\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "\n",
    "img = Image.open(\"/tmp/giant_panda.jpg\")\n",
    "# Preprocess image\n",
    "image_size = 224\n",
    "tfms = transforms.Compose([transforms.Resize(image_size), \n",
    "                           transforms.CenterCrop(image_size), \n",
    "                           transforms.ToTensor(),\n",
    "                           transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),])\n",
    "x = tfms(img).unsqueeze(0)\n",
    "\n",
    "plt.imshow(img)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "l5xn0U7bsXp6"
   },
   "source": [
    "# Classify\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    y_pred = model(x)\n",
    "\n",
    "# Print predictions\n",
    "print('-----')\n",
    "for idx in torch.topk(y_pred, k=5).indices.squeeze(0).tolist():\n",
    "    prob = torch.softmax(y_pred, dim=1)[0, idx].item()\n",
    "    print('{label:<75} ({p:.2f}%)'.format(label=labels[str(idx)], p=prob*100))"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6p_g9Z4KsXp6"
   },
   "source": [
    "## Dataflow\n",
    "\n",
    "Let's setup the dataflow:\n",
    "- load CIFAR100 train and test datasets\n",
    "- setup train/test image transforms\n",
    "- setup train/test data loaders\n",
    "\n",
    "According to the paper authors borrowed training settings from other publications and the dataflow for CIFAR100 is the following:\n",
    "\n",
    "- input images to the network during training are resized to 224x224\n",
    "- horizontally flipped randomly and augmented using cutout.\n",
    "- each mini-batch contained 256 examples\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "Xm1nwsuVsXp7"
   },
   "source": [
    "from torchvision.datasets.cifar import CIFAR100 \n",
    "from torchvision.transforms import Compose, RandomCrop, Pad, RandomHorizontalFlip, Resize\n",
    "from torchvision.transforms import ToTensor, Normalize\n",
    "\n",
    "from torch.utils.data import Subset"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "t4dcDFtxsXp7"
   },
   "source": [
    "path = \"/tmp/cifar100\"\n",
    "\n",
    "from PIL.Image import BICUBIC\n",
    "\n",
    "\n",
    "train_transform = Compose([\n",
    "    Resize(256, BICUBIC),\n",
    "    RandomCrop(224),\n",
    "    RandomHorizontalFlip(),\n",
    "    ToTensor(),\n",
    "    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "])\n",
    "\n",
    "test_transform = Compose([\n",
    "    Resize(224, BICUBIC),    \n",
    "    ToTensor(),\n",
    "    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "])\n",
    "\n",
    "\n",
    "train_dataset = CIFAR100(root=path, train=True, transform=train_transform, download=True)\n",
    "test_dataset = CIFAR100(root=path, train=False, transform=test_transform, download=False)\n",
    "\n",
    "train_eval_indices = [random.randint(0, len(train_dataset) - 1) for i in range(len(test_dataset))]\n",
    "train_eval_dataset = Subset(train_dataset, train_eval_indices)\n",
    "\n",
    "\n",
    "len(train_dataset), len(test_dataset), len(train_eval_dataset)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "pwOEeRFOsXp7"
   },
   "source": [
    "from torch.utils.data import DataLoader\n",
    "\n",
    "\n",
    "batch_size = 172\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size, num_workers=20, \n",
    "                          shuffle=True, drop_last=True, pin_memory=True)\n",
    "\n",
    "test_loader = DataLoader(test_dataset, batch_size=batch_size, num_workers=20, \n",
    "                         shuffle=False, drop_last=False, pin_memory=True)\n",
    "\n",
    "eval_train_loader = DataLoader(train_eval_dataset, batch_size=batch_size, num_workers=20, \n",
    "                               shuffle=False, drop_last=False, pin_memory=True)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "UXCG-s1JsXp7"
   },
   "source": [
    "import torchvision.utils as vutils\n",
    "\n",
    "# Plot some training images\n",
    "batch = next(iter(train_loader))\n",
    "\n",
    "plt.figure(figsize=(16, 8))\n",
    "plt.axis(\"off\")\n",
    "plt.title(\"Training Images\")\n",
    "plt.imshow( \n",
    "    vutils.make_grid(batch[0][:16], padding=2, normalize=True).cpu().numpy().transpose((1, 2, 0))\n",
    ")\n",
    "\n",
    "batch = None\n",
    "torch.cuda.empty_cache()"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0ty_zhLUsXp8"
   },
   "source": [
    "## Finetunning model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dvVxC-_TsXp9"
   },
   "source": [
    "As we are interested to finetune the model to CIFAR-100, we will replace the classification fully-connected layer (ImageNet-1000 vs CIFAR-100)."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "O7D5BIlxsXp9"
   },
   "source": [
    "model.head[6].in_features, model.head[6].out_features"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "mwLQro9csXp-"
   },
   "source": [
    "model.head[6] = nn.Linear(1280, 100)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "1tcFlIpGsXp-"
   },
   "source": [
    "model.head[6].in_features, model.head[6].out_features"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3uS2NMt8sXp_"
   },
   "source": [
    "We will finetune the model on GPU with AMP fp32/fp16 using nvidia/apex package."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "Fv8lR4HlsXp_"
   },
   "source": [
    "assert torch.cuda.is_available()\n",
    "assert torch.backends.cudnn.enabled, \"NVIDIA/Apex:Amp requires cudnn backend to be enabled.\"\n",
    "torch.backends.cudnn.benchmark = True\n",
    "\n",
    "device = \"cuda\""
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "TUhjJfLTsXp_"
   },
   "source": [
    "model = model.to(device)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2YB8-FGXsXqB"
   },
   "source": [
    "Let's setup cross-entropy as criterion and SGD as optimizer.\n",
    "\n",
    "We will split model parameters into 2 groups: \n",
    "\n",
    "    1) feature extractor (pretrained weights)\n",
    "    2) classifier (random weights)\n",
    "\n",
    "and define different learning rates for these groups (via learning rate scheduler)."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "TGxSb_MOsXqB"
   },
   "source": [
    "from itertools import chain\n",
    "\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "lr = 0.01\n",
    "\n",
    "optimizer = optim.SGD([\n",
    "    {\n",
    "        \"params\": chain(model.stem.parameters(), model.blocks.parameters()),\n",
    "        \"lr\": lr * 0.1,\n",
    "    },\n",
    "    {\n",
    "        \"params\": model.head[:6].parameters(),\n",
    "        \"lr\": lr * 0.2,\n",
    "    },    \n",
    "    {\n",
    "        \"params\": model.head[6].parameters(), \n",
    "        \"lr\": lr\n",
    "    }], \n",
    "    momentum=0.9, weight_decay=0.001, nesterov=True)\n"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "V6stcy6ksXqB"
   },
   "source": [
    "from torch.optim.lr_scheduler import ExponentialLR\n",
    "\n",
    "lr_scheduler = ExponentialLR(optimizer, gamma=0.975)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "yPuwNYeZsXqB"
   },
   "source": [
    "try:\n",
    "    from apex import amp\n",
    "except ImportError:\n",
    "    raise ImportError(\"Please install apex from https://www.github.com/nvidia/apex to run this example.\")\n",
    "\n",
    "\n",
    "# Initialize Amp\n",
    "model, optimizer = amp.initialize(model, optimizer, opt_level=\"O2\", num_losses=1)\n"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5UtGQtQFsXqC"
   },
   "source": [
    "Next, let's define a single iteration function `update_fn`. This function is then used by `ignite.engine.Engine` to update model while running over the input data."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "MjMUQITRsXqC"
   },
   "source": [
    "from ignite.utils import convert_tensor\n",
    "\n",
    "\n",
    "def update_fn(engine, batch):\n",
    "    model.train()\n",
    "\n",
    "    x = convert_tensor(batch[0], device=device, non_blocking=True)\n",
    "    y = convert_tensor(batch[1], device=device, non_blocking=True)\n",
    "    \n",
    "    y_pred = model(x)\n",
    "    \n",
    "    # Compute loss \n",
    "    loss = criterion(y_pred, y)    \n",
    "\n",
    "    with amp.scale_loss(loss, optimizer) as scaled_loss:\n",
    "        scaled_loss.backward()\n",
    "\n",
    "    optimizer.step()\n",
    "    \n",
    "    return {\n",
    "        \"batchloss\": loss.item(),\n",
    "    }    "
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "E772HMiasXqC"
   },
   "source": [
    "Let's check `update_fn`"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "MLFVuQaDsXqD"
   },
   "source": [
    "batch = next(iter(train_loader))\n",
    "\n",
    "res = update_fn(engine=None, batch=batch)\n",
    "\n",
    "batch = None\n",
    "torch.cuda.empty_cache()\n",
    "\n",
    "res"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ibK059y1sXqD"
   },
   "source": [
    "Now let's define a trainer and add some practical handlers:\n",
    "- log to tensorboard: losses, metrics, lr\n",
    "- progress bar\n",
    "- models/optimizers checkpointing"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "rqrmwVuTsXqD"
   },
   "source": [
    "from ignite.engine import Engine, Events, create_supervised_evaluator\n",
    "from ignite.metrics import RunningAverage, Accuracy, Precision, Recall, Loss, TopKCategoricalAccuracy\n",
    "\n",
    "from ignite.handlers import TensorboardLogger\n",
    "from ignite.handlers.tensorboard_logger import OutputHandler, OptimizerParamsHandler"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "B2JI5oXosXqD"
   },
   "source": [
    "trainer = Engine(update_fn)\n",
    "\n",
    "\n",
    "def output_transform(out):\n",
    "    return out['batchloss']\n",
    "\n",
    "\n",
    "RunningAverage(output_transform=output_transform).attach(trainer, \"batchloss\")"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "2_WWUKuHsXqD"
   },
   "source": [
    "from datetime import datetime\n",
    "\n",
    "exp_name = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
    "log_path = \"/tmp/finetune_efficientnet_cifar100/{}\".format(exp_name)\n",
    "tb_logger = TensorboardLogger(log_dir=log_path)\n",
    "\n",
    "\n",
    "tb_logger.attach(trainer, \n",
    "                 log_handler=OutputHandler('training', ['batchloss', ]), \n",
    "                 event_name=Events.ITERATION_COMPLETED)\n",
    "\n",
    "print(\"Experiment name: \", exp_name)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "R6BPVSZqsXqE"
   },
   "source": [
    "Let's setup learning rate scheduling:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ZVHrbB5esXqE"
   },
   "source": [
    "trainer.add_event_handler(Events.EPOCH_COMPLETED, lambda engine: lr_scheduler.step())\n",
    "\n",
    "# Log optimizer parameters\n",
    "tb_logger.attach(trainer,\n",
    "                 log_handler=OptimizerParamsHandler(optimizer, \"lr\"), \n",
    "                 event_name=Events.EPOCH_STARTED)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "Bp_hy7Y0sXqE"
   },
   "source": [
    "from ignite.handlers import ProgressBar\n",
    "\n",
    "# Iteration-wise progress bar\n",
    "# ProgressBar(bar_format=\"\").attach(trainer, metric_names=['batchloss',])\n",
    "\n",
    "# Epoch-wise progress bar with display of training losses\n",
    "ProgressBar(persist=True, bar_format=\"\").attach(trainer, \n",
    "                                                event_name=Events.EPOCH_STARTED, \n",
    "                                                closing_event_name=Events.COMPLETED)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EM6NNDh8sXqE"
   },
   "source": [
    "Let's create two evaluators to compute metrics on train/test images and log them to Tensorboard:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "yApMSd47sXqF"
   },
   "source": [
    "metrics = {\n",
    "    'Loss': Loss(criterion),\n",
    "    'Accuracy': Accuracy(),\n",
    "    'Precision': Precision(average=True),\n",
    "    'Recall': Recall(average=True),\n",
    "    'Top-5 Accuracy': TopKCategoricalAccuracy(k=5)\n",
    "}\n",
    "\n",
    "\n",
    "evaluator = create_supervised_evaluator(model, metrics=metrics, \n",
    "                                        device=device, non_blocking=True)\n",
    "\n",
    "train_evaluator = create_supervised_evaluator(model, metrics=metrics, \n",
    "                                              device=device, non_blocking=True)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "_Xn8INiCsXqF"
   },
   "source": [
    "from ignite.handlers import global_step_from_engine\n",
    "\n",
    "\n",
    "def run_evaluation(engine):\n",
    "    train_evaluator.run(eval_train_loader)\n",
    "    evaluator.run(test_loader)\n",
    "\n",
    "\n",
    "trainer.add_event_handler(Events.EPOCH_STARTED(every=3), run_evaluation)\n",
    "trainer.add_event_handler(Events.COMPLETED, run_evaluation)\n",
    "\n",
    "\n",
    "# Log train eval metrics:\n",
    "tb_logger.attach_output_handler(\n",
    "    train_evaluator,\n",
    "    event_name=Events.EPOCH_COMPLETED,\n",
    "    tag=\"training\",\n",
    "    metric_names=list(metrics.keys()),\n",
    "    global_step_transform=global_step_from_engine(trainer)\n",
    ")\n",
    "\n",
    "# Log val metrics:\n",
    "tb_logger.attach_output_handler(\n",
    "    evaluator,\n",
    "    event_name=Events.EPOCH_COMPLETED,\n",
    "    tag=\"test\",\n",
    "    metric_names=list(metrics.keys()),\n",
    "    global_step_transform=global_step_from_engine(trainer)\n",
    ")"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Sx_J-SHisXqG"
   },
   "source": [
    "Now let's setup the best model checkpointing, early stopping:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "tC5UBpXfsXqG"
   },
   "source": [
    "import logging\n",
    "\n",
    "# Setup engine &  logger\n",
    "def setup_logger(logger):\n",
    "    handler = logging.StreamHandler()\n",
    "    formatter = logging.Formatter(\"%(asctime)s %(name)-12s %(levelname)-8s %(message)s\")\n",
    "    handler.setFormatter(formatter)\n",
    "    logger.addHandler(handler)\n",
    "    logger.setLevel(logging.INFO)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "LKXnwpxBsXqG"
   },
   "source": [
    "from ignite.handlers import Checkpoint, DiskSaver, EarlyStopping, TerminateOnNan\n",
    "\n",
    "\n",
    "trainer.add_event_handler(Events.ITERATION_COMPLETED, TerminateOnNan())\n",
    "\n",
    "\n",
    "# Store the best model\n",
    "def default_score_fn(engine):\n",
    "    score = engine.state.metrics['Accuracy']\n",
    "    return score\n",
    "\n",
    "# Force filename to model.pt to ease the rerun of the notebook\n",
    "disk_saver = DiskSaver(dirname=log_path)\n",
    "best_model_handler = Checkpoint(to_save={'model': model}, \n",
    "                                save_handler=disk_saver, \n",
    "                                filename_pattern=\"{name}.{ext}\", \n",
    "                                n_saved=1)\n",
    "evaluator.add_event_handler(Events.COMPLETED, best_model_handler)\n",
    "\n",
    "# Add early stopping\n",
    "es_patience = 10\n",
    "es_handler = EarlyStopping(patience=es_patience, score_function=default_score_fn, trainer=trainer)\n",
    "evaluator.add_event_handler(Events.COMPLETED, es_handler)\n",
    "setup_logger(es_handler.logger)\n",
    "\n",
    "\n",
    "# Clear cuda cache between training/testing\n",
    "def empty_cuda_cache(engine):\n",
    "    torch.cuda.empty_cache()\n",
    "    import gc\n",
    "    gc.collect()\n",
    "\n",
    "\n",
    "trainer.add_event_handler(Events.EPOCH_COMPLETED, empty_cuda_cache)\n",
    "evaluator.add_event_handler(Events.COMPLETED, empty_cuda_cache)\n",
    "train_evaluator.add_event_handler(Events.COMPLETED, empty_cuda_cache)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "LBP3QVQXsXqG"
   },
   "source": [
    "num_epochs = 100\n",
    "\n",
    "\n",
    "trainer.run(train_loader, max_epochs=num_epochs)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gsXB45JhsXqH"
   },
   "source": [
    "Finetunning results:\n",
    "\n",
    "- Test dataset:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "mX-t-Sq6sXqH"
   },
   "source": [
    "evaluator.state.metrics"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ax_YKq4isXqH"
   },
   "source": [
    "- Training subset:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "i4CA6YcRsXqI"
   },
   "source": [
    "train_evaluator.state.metrics"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ghxtoGLVsXqI"
   },
   "source": [
    "Obviously, our training settings is not the optimal one and the delta between our result and the paper's one is about 5%."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "C-7ccZ0QsXqJ"
   },
   "source": [
    "## Inference\n",
    "\n",
    "Let's load the best model and recompute evaluation metrics on test dataset with a very basic Test-Time-Augmentation to boost the performances.\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "x4BU5IOIsXqK"
   },
   "source": [
    "best_model = EfficientNet()\n",
    "best_model.load_state_dict(torch.load(log_path + \"/model.pt\"))"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "co3E6dcesXqK"
   },
   "source": [
    "metrics = {\n",
    "    'Accuracy': Accuracy(),\n",
    "    'Precision': Precision(average=True),\n",
    "    'Recall': Recall(average=True),\n",
    "}\n",
    "\n",
    "\n",
    "def inference_update_with_tta(engine, batch):\n",
    "    best_model.eval()\n",
    "    with torch.no_grad():\n",
    "        x, y = batch        \n",
    "        # Let's compute final prediction as a mean of predictions on x and flipped x\n",
    "        y_pred1 = best_model(x)\n",
    "        y_pred2 = best_model(x.flip(dims=(-1, )))\n",
    "        y_pred = 0.5 * (y_pred1 + y_pred2)\n",
    "\n",
    "        return y_pred, y\n",
    "\n",
    "\n",
    "inferencer = Engine(inference_update_with_tta)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "u68ntSHDsXqK"
   },
   "source": [
    "for name, metric in metrics.items():\n",
    "    metric.attach(inferencer, name)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "W3YDadJ5sXqK"
   },
   "source": [
    "ProgressBar(desc=\"Inference\").attach(inferencer)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "4fQTu7DOsXqK"
   },
   "source": [
    "result_state = inferencer.run(test_loader, max_epochs=1)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1vlgsopvsXqL"
   },
   "source": [
    "Finally, we obtain similar scores:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "F6vCp6lpsXqL"
   },
   "source": [
    "result_state.metrics"
   ],
   "execution_count": null,
   "outputs": []
  }
 ]
}