Imagenette_v2.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Imagenette_v2.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/davidcpage/Imagenette-experiments/blob/master/Imagenette_v2.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q8mLtXI5gfGV",
        "colab_type": "text"
      },
      "source": [
        "The aim of this notebook is to convert the Imagenette/woof training examples from https://github.com/lessw2020/Ranger-Mish-ImageWoof-5 to the new fastai v2 codebase. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FU0TYajRTt23",
        "colab_type": "text"
      },
      "source": [
        "### Setup"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tZSzbQ9Ag47l",
        "colab_type": "text"
      },
      "source": [
        "Install fastai2. You may need to restart after installing to pick up new versions."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vLGuUwFpLzc1",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!python -m pip install typeguard\n",
        "!python -m pip install --upgrade pillow fastprogress\n",
        "!python -m pip install git+https://github.com/fastai/fastai2\n",
        "\n",
        "RANGER = 'https://raw.githubusercontent.com/lessw2020/Ranger-Mish-ImageWoof-5/master/ranger.py'\n",
        "MXRESNET = 'https://raw.githubusercontent.com/lessw2020/Ranger-Mish-ImageWoof-5/master/mxresnet.py'\n",
        "UTILS = 'https://raw.githubusercontent.com/davidcpage/Imagenette-experiments/master/utils.py'\n",
        "\n",
        "!wget $RANGER -O ranger.py\n",
        "!wget $MXRESNET -O mxresnet.py\n",
        "!wget $UTILS -O utils.py"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k0Eqi9AJotZP",
        "colab_type": "text"
      },
      "source": [
        "Install Nvidia DALI, which we use for fast dataloading/augmentation below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vuoBhXAcm0Sj",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!python -m pip install --extra-index-url https://developer.download.nvidia.com/compute/redist/cuda/10.0 nvidia-dali"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i7c9ywHZhNiT",
        "colab_type": "text"
      },
      "source": [
        "Basic imports and device setup."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FbkdYs8IM37W",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from functools import partial\n",
        "import numpy as np\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "device = torch.device(torch.cuda.current_device())\n",
        "torch.backends.cudnn.benchmark = True"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BvQYKhXLhr8I",
        "colab_type": "text"
      },
      "source": [
        "#### Params"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_wO66U5HPRvI",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "size = 128\n",
        "bs = 64\n",
        "\n",
        "random_aspect_ratio = (3/4, 4/3)\n",
        "random_area = (0.35, 1.)\n",
        "val_xtra_size = 32 \n",
        "interpolation = 2 #FIXME: what is this?"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QFzsy_gcNlxU",
        "colab_type": "text"
      },
      "source": [
        "### Fastai v1 training"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bRe8tInlNvmR",
        "colab_type": "text"
      },
      "source": [
        "First let's establish a baseline on the v1 codebase. The code here is essentially taken from https://github.com/lessw2020/Ranger-Mish-ImageWoof-5/blob/master/train.py. \n",
        "\n",
        "**NB:** throughout the notebook we use fully qualified names for imported functions to avoid name-clashes between the two fastai codebases."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bFpzkgp4PGz_",
        "colab_type": "code",
        "outputId": "e227b827-b251-47a5-ac79-7df6be69b393",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "import fastai.vision, fastai.vision.data, fastai.basic_train\n",
        "import fastai.metrics, fastai.layers, fastai.callbacks\n",
        "\n",
        "import ranger\n",
        "import mxresnet\n",
        "\n",
        "data_dir = fastai.datasets.untar_data(fastai.datasets.URLs.IMAGENETTE_320)"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Mish activation loaded...\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DbwxPTvcRgG5",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "data = lambda data_dir=data_dir: (\n",
        "    fastai.vision.data.ImageList.from_folder(data_dir)\n",
        "            .split_by_folder(valid='val')\n",
        "            .label_from_folder().transform(([fastai.vision.flip_lr(p=0.5)], []), size=size)\n",
        "            .databunch(bs=bs, num_workers=4)\n",
        "            .presize(size, scale=random_area, ratio=random_aspect_ratio, val_xtra_size=val_xtra_size, interpolation=interpolation)\n",
        "            .normalize(fastai.vision.imagenet_stats)\n",
        "    )\n",
        "\n",
        "def flat_then_cosine_sched(learn, n_batch, lr, annealing_start):\n",
        "    return fastai.callbacks.GeneralScheduler(learn, phases=[\n",
        "        fastai.callbacks.TrainingPhase(annealing_start*n_batch).schedule_hp('lr', lr),\n",
        "        fastai.callbacks.TrainingPhase((1-annealing_start)*n_batch).schedule_hp('lr', lr, anneal=fastai.callback.annealing_cos)     \n",
        "    ])\n",
        "\n",
        "def fit_flat_cos(learn, num_epoch, lr=4e-3, annealing_start=0.72):\n",
        "    learn.fit(num_epoch, callbacks=[\n",
        "        flat_then_cosine_sched(learn, len(learn.data.train_dl) * num_epoch, lr=lr, annealing_start=annealing_start)])\n",
        "    return learn\n",
        "\n",
        "learner = partial(fastai.basic_train.Learner, \n",
        "                  wd=1e-2,\n",
        "                  opt_func=partial(ranger.Ranger, betas=(0.95, 0.99), eps=1e-6),\n",
        "                  metrics=(fastai.metrics.accuracy,),\n",
        "                  bn_wd=False, true_wd=True,\n",
        "                  loss_func=fastai.layers.LabelSmoothingCrossEntropy(),\n",
        "                  )"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q6s8x8vriI4Q",
        "colab_type": "text"
      },
      "source": [
        "In order to keep things relatively fast, let's train an mxresnet18 model for 5 epochs:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gOGmIScpAmLw",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 204
        },
        "outputId": "71f9079f-cbed-4d72-c7c2-fe067f111115"
      },
      "source": [
        "model = mxresnet.mxresnet18(c_out=10, sa=1, sym=0)\n",
        "learn = fit_flat_cos(learner(data(), model).to_fp16(), num_epoch=5)"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: left;\">\n",
              "      <th>epoch</th>\n",
              "      <th>train_loss</th>\n",
              "      <th>valid_loss</th>\n",
              "      <th>accuracy</th>\n",
              "      <th>time</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <td>0</td>\n",
              "      <td>1.441207</td>\n",
              "      <td>1.277145</td>\n",
              "      <td>0.684000</td>\n",
              "      <td>00:44</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>1</td>\n",
              "      <td>1.221641</td>\n",
              "      <td>1.014491</td>\n",
              "      <td>0.808000</td>\n",
              "      <td>00:44</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>2</td>\n",
              "      <td>1.137839</td>\n",
              "      <td>1.073034</td>\n",
              "      <td>0.772000</td>\n",
              "      <td>00:44</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>3</td>\n",
              "      <td>1.059366</td>\n",
              "      <td>0.949235</td>\n",
              "      <td>0.842000</td>\n",
              "      <td>00:43</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>4</td>\n",
              "      <td>0.946012</td>\n",
              "      <td>0.840164</td>\n",
              "      <td>0.884000</td>\n",
              "      <td>00:43</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2ASahDA1VOw8",
        "colab_type": "text"
      },
      "source": [
        "### Nvidia DALI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aPojLp2SjYLO",
        "colab_type": "text"
      },
      "source": [
        "We can speed things up (especially for small models) by using Nvidia DALI to do the dataloading/augmentation. For now we will use this for both fastai v1 and v2 models although we will want to test v2 dataloading at the end. The details of the code are not very interesting.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2Sq6NHHIVQpE",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import nvidia.dali.ops as ops\n",
        "import nvidia.dali.types as types\n",
        "\n",
        "imagenet_stats = ((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))\n",
        "\n",
        "def imagenet_train_graph(data_dir, size, random_aspect_ratio, random_area, \n",
        "                interp_type=types.INTERP_TRIANGULAR,\n",
        "                stats=imagenet_stats):\n",
        "    inputs = ops.FileReader(file_root=data_dir, random_shuffle=True)\n",
        "    decode = ops.ImageDecoderRandomCrop(device='mixed',\n",
        "            random_aspect_ratio=random_aspect_ratio, random_area=random_area)\n",
        "    resize = ops.Resize(device='gpu', resize_x=size, resize_y=size, \n",
        "                        interp_type=interp_type)\n",
        "    mean, std = [[x*255 for x in stat] for stat in stats]\n",
        "    crop_mirror_norm = ops.CropMirrorNormalize(\n",
        "                        device='gpu', output_dtype=types.FLOAT16, \n",
        "                        crop=(size, size), mean=mean, std=std)\n",
        "    coin = ops.CoinFlip(probability=0.5)\n",
        "\n",
        "    def define_graph():    \n",
        "        jpegs, labels = inputs(name='Reader')\n",
        "        output = crop_mirror_norm(resize(decode(jpegs)), mirror=coin())\n",
        "        return [output, labels]\n",
        "    return define_graph\n",
        "\n",
        "def imagenet_valid_graph(data_dir, size, val_xtra_size, mirror=0,\n",
        "                interp_type=types.INTERP_TRIANGULAR, \n",
        "                stats=imagenet_stats):\n",
        "    inputs = ops.FileReader(file_root=data_dir, random_shuffle=False)\n",
        "    decode = ops.ImageDecoder(device='mixed', output_type=types.RGB)\n",
        "    resize = ops.Resize(device='gpu', resize_shorter=size+val_xtra_size, \n",
        "                        interp_type=interp_type)\n",
        "    mean, std = [[x*255 for x in stat] for stat in stats]\n",
        "    crop_mirror_norm = ops.CropMirrorNormalize(\n",
        "                        device='gpu', output_dtype=types.FLOAT16,\n",
        "                        crop=(size, size), mean=mean, std=std, mirror=mirror)\n",
        "    \n",
        "    def define_graph():\n",
        "        jpegs, labels = inputs(name='Reader')\n",
        "        output = crop_mirror_norm(resize(decode(jpegs)))\n",
        "        return [output, labels]\n",
        "    return define_graph"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "km5FiFyDkmwJ",
        "colab_type": "text"
      },
      "source": [
        "The validation set of Imagenette is tiny, consisting of 500 examples only. This leads to substantial noise in the validation accuracies. To help a little, we are going to concatenate a left-right flipped version onto to the validation set to get an effective 1000 examples:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "B1ZFAN2lVQjr",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from utils import DALIDataLoader, Chain, MockV1DataBunch\n",
        "\n",
        "train_dl = lambda folder, bs, seed=-1: (\n",
        "        DALIDataLoader(imagenet_train_graph(folder, size, random_aspect_ratio, random_area), bs, drop_last=True, device=device, seed=seed))\n",
        "valid_dl = lambda folder, bs, : Chain(\n",
        "        DALIDataLoader(imagenet_valid_graph(folder, size, val_xtra_size), bs, drop_last=False, device=device),\n",
        "        DALIDataLoader(imagenet_valid_graph(folder, size, val_xtra_size, mirror=1), bs, drop_last=False, device=device),\n",
        "    )\n",
        "\n",
        "dali_data = lambda data_dir=data_dir, bs=bs: MockV1DataBunch(train_dl(data_dir/'train', bs), valid_dl(data_dir/'val', bs))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RKV_fGfmk6f2",
        "colab_type": "text"
      },
      "source": [
        "Now we are ready to test training on the v1 codebase with DALI dataloading:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wQGkC6nCDGkw",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 204
        },
        "outputId": "2969dc4d-0620-491b-d0ab-e37cd005a6c2"
      },
      "source": [
        "model = mxresnet.mxresnet18(c_out=10, sa=1, sym=0)\n",
        "learn = fit_flat_cos(learner(dali_data(), model).to_fp16(), num_epoch=5)"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: left;\">\n",
              "      <th>epoch</th>\n",
              "      <th>train_loss</th>\n",
              "      <th>valid_loss</th>\n",
              "      <th>accuracy</th>\n",
              "      <th>time</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <td>0</td>\n",
              "      <td>1.422536</td>\n",
              "      <td>1.233743</td>\n",
              "      <td>0.716000</td>\n",
              "      <td>00:18</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>1</td>\n",
              "      <td>1.202731</td>\n",
              "      <td>1.047867</td>\n",
              "      <td>0.811000</td>\n",
              "      <td>00:18</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>2</td>\n",
              "      <td>1.094656</td>\n",
              "      <td>0.991178</td>\n",
              "      <td>0.813000</td>\n",
              "      <td>00:18</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>3</td>\n",
              "      <td>1.026366</td>\n",
              "      <td>0.925935</td>\n",
              "      <td>0.843000</td>\n",
              "      <td>00:18</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>4</td>\n",
              "      <td>0.917531</td>\n",
              "      <td>0.842249</td>\n",
              "      <td>0.877000</td>\n",
              "      <td>00:18</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UTVevmh6lyPF",
        "colab_type": "text"
      },
      "source": [
        "The training is substantially faster (at least on a GCP T4 GPU with 4 cpu cores) and accuracy is similar (although noise is still very large.)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Um6YrUVZkNpA",
        "colab_type": "text"
      },
      "source": [
        "### Fastai v2 model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w_cwvCkprxH5",
        "colab_type": "text"
      },
      "source": [
        "Next we want to compare a model using the v2 codebase to the v1 model above. Here is the v1 model again:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AjMFCkw2r8qf",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model_v1 = partial(mxresnet.mxresnet18, c_out=10, sa=1, sym=0)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CCa6wKrZr_M8",
        "colab_type": "text"
      },
      "source": [
        "A similar model is available in v2. To start with let's use the Mish activation class from the v1 model to minimise differences."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WN-bLy6skieq",
        "colab_type": "code",
        "outputId": "c2c9bf1e-20d6-4008-9554-f159a3e4bff2",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "import fastai2.vision.models.xresnet\n",
        "mish = mxresnet.Mish()\n",
        "model_v2 = partial(fastai2.vision.models.xresnet.xresnet18, c_out=10, sa=1, sym=0, act_cls=(lambda: mish))"
      ],
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Mish activation loaded...\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qTYAAEhqyiqg",
        "colab_type": "text"
      },
      "source": [
        "First let's compare the types of modules in the two models:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "688mhsD5EAUs",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 187
        },
        "outputId": "36bded59-1f63-469d-da76-abeced0fe7b6"
      },
      "source": [
        "s1 = set(type(x) for x in model_v1().modules())\n",
        "s2 = set(type(x) for x in model_v2().modules())\n",
        "s1^s2 #symmetric difference"
      ],
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{fastai2.layers.ConvLayer,\n",
              " fastai2.layers.Flatten,\n",
              " fastai2.layers.ResBlock,\n",
              " fastai2.vision.models.xresnet.XResNet,\n",
              " mxresnet.Flatten,\n",
              " mxresnet.MXResNet,\n",
              " mxresnet.ResBlock,\n",
              " mxresnet.SimpleSelfAttention,\n",
              " torch.nn.modules.activation.ReLU,\n",
              " torch.nn.modules.conv.Conv1d}"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 11
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rf9GH9w76jGv",
        "colab_type": "text"
      },
      "source": [
        "Some of these are implementation specific compound layers. Let's ignore them for now and we will come back to them later if necessary."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fyoqGP6lEKST",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "types_to_ignore = {\n",
        "    torch.nn.Sequential,\n",
        "\n",
        "    mxresnet.Flatten, \n",
        "    mxresnet.MXResNet, \n",
        "    mxresnet.ResBlock, \n",
        "\n",
        "    fastai2.layers.Flatten, \n",
        "    fastai2.vision.models.xresnet.XResNet,\n",
        "    fastai2.layers.ResBlock, \n",
        "    fastai2.layers.ConvLayer, \n",
        "}\n",
        "\n",
        "filtered_modules = lambda model: (x for x in model.modules() if type(x) not in types_to_ignore)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uuM3Jzit-8oy",
        "colab_type": "text"
      },
      "source": [
        "Now let's compare the filtered modules in a bit more detail by comparing `repr` strings."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "O_uvvNObEPQa",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 85
        },
        "outputId": "42c7e82e-278c-4dac-87eb-787136c7c7e9"
      },
      "source": [
        "s1 = set(repr(x) for x in filtered_modules(model_v1()))\n",
        "s2 = set(repr(x) for x in filtered_modules(model_v2()))\n",
        "s1 ^ s2"
      ],
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'Conv1d(64, 64, kernel_size=(1,), stride=(1,), bias=False)',\n",
              " 'Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)',\n",
              " 'ReLU()',\n",
              " 'SimpleSelfAttention(\\n  (conv): Conv1d(64, 64, kernel_size=(1,), stride=(1,), bias=False)\\n)'}"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 13
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ysFrMZH-_mTU",
        "colab_type": "text"
      },
      "source": [
        "Let's locate the offending modules in the v1 model:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2jZCkibqERhG",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 68
        },
        "outputId": "2ca86631-4d8a-461c-f17f-ae7747319b88"
      },
      "source": [
        "{k: x for k,x in model_v1().named_modules() if repr(x) in s1^s2}"
      ],
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'4.1.sa': SimpleSelfAttention(\n",
              "   (conv): Conv1d(64, 64, kernel_size=(1,), stride=(1,), bias=False)\n",
              " ), '4.1.sa.conv': Conv1d(64, 64, kernel_size=(1,), stride=(1,), bias=False)}"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 14
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5AMO0MgE_q4C",
        "colab_type": "text"
      },
      "source": [
        "and in the v2 model:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "j4IXMDMwETBp",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 85
        },
        "outputId": "a8e49cc0-4e34-4933-e3f9-243e04978ec4"
      },
      "source": [
        "{k: x for k,x in model_v2().named_modules() if repr(x) in s1^s2}"
      ],
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'0.2': ReLU(),\n",
              " '1.0': Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),\n",
              " '1.2': ReLU(),\n",
              " '2.2': ReLU()}"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 15
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0yuhTB3pACPX",
        "colab_type": "text"
      },
      "source": [
        "The v1 model has a SimpleSelfAttention layer which is missing from the v2 model, whilst the v2 model still has some ReLU activations at early layers and one Conv layer has a different shape. Let's fix these issues with a modified v2 model:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KFdBgDoWEbIi",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class XResNet(nn.Sequential):\n",
        "    def __init__(self, expansion, layers, c_in=3, c_out=1000, \n",
        "                 sa=False, sym=False, act_cls=fastai2.basics.defaults.activation,\n",
        "                 ):\n",
        "        stem = []\n",
        "        sizes = [c_in, 16,32,64] if c_in < 3 else [c_in, 32, 64, 64] \n",
        "        for i in range(3):\n",
        "            stem.append(fastai2.layers.ConvLayer(sizes[i], sizes[i+1], stride=2 if i==0 else 1, act_cls=act_cls))\n",
        "\n",
        "        block_szs = [64//expansion,64,128,256,512] +[256]*(len(layers)-4)\n",
        "        blocks = [self._make_layer(expansion, ni=block_szs[i], nf=block_szs[i+1], blocks=l, stride=1 if i==0 else 2,\n",
        "                                  sa=sa if i==len(layers)-4 else False, sym=sym, act_cls=act_cls)\n",
        "                  for i,l in enumerate(layers)]\n",
        "        super().__init__(\n",
        "            *stem,\n",
        "            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),\n",
        "            *blocks,\n",
        "            nn.AdaptiveAvgPool2d(1), fastai2.layers.Flatten(),\n",
        "            nn.Linear(block_szs[-1]*expansion, c_out),\n",
        "        )\n",
        "        fastai2.vision.models.xresnet.init_cnn(self)\n",
        "\n",
        "    def _make_layer(self, expansion, ni, nf, blocks, stride, sa, sym, act_cls):\n",
        "        return nn.Sequential(\n",
        "            *[fastai2.layers.ResBlock(expansion, ni if i==0 else nf, nf, stride if i==0 else 1,\n",
        "                      sa if i==(blocks-1) else False, sym=sym, act_cls=act_cls)\n",
        "              for i in range(blocks)])\n",
        "        \n",
        "xresnet18 = partial(XResNet, expansion=1, layers=[2,2,2,2])\n",
        "xresnet50 = partial(XResNet, expansion=4, layers=[3,4,6,3])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AXj4chBgB-OE",
        "colab_type": "text"
      },
      "source": [
        "Let's instantiate the new model and run our check from before to compare with the v1 model:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hiXKnEKT_tJN",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model_v2b = partial(xresnet18, c_out=10, sa=1, sym=0, act_cls=(lambda: mish))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OY8OQr3y_tEQ",
        "colab_type": "code",
        "outputId": "7a9246e7-13d3-47ca-88a8-b52d674c8dc6",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "types_to_ignore.add(XResNet)\n",
        "\n",
        "s1 = set(repr(x) for x in filtered_modules(model_v1()))\n",
        "s2 = set(repr(x) for x in filtered_modules(model_v2b()))\n",
        "s1 ^ s2"
      ],
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "set()"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 18
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d2-cKU7MCrUn",
        "colab_type": "text"
      },
      "source": [
        "Great. Next we'd like to check that the forward computation of the two models ties out. The difficulty with this is that initialisation calls random number generators and even if we fix random seeds beforehand, any difference in the sequence of random calls for the two models will lead to divergence.\n",
        "\n",
        "We can check that other details of the forward computation agree by attempting to set the same initialisation for both models:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7uw_sRAxEXK7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def reset_dummy_init(model, seed=123):\n",
        "    for m in filtered_modules(model):\n",
        "        if hasattr(m, 'reset_parameters'):\n",
        "            torch.manual_seed(seed)\n",
        "            m.reset_parameters()\n",
        "    return model\n",
        "\n",
        "def compare_fwd(model1, model2):\n",
        "    random_batch = torch.randn(bs,3,size,size, device=device)\n",
        "    assert np.allclose(\n",
        "        model1.to(device)(random_batch).detach().cpu().numpy(), \n",
        "        model2.to(device)(random_batch).detach().cpu().numpy()\n",
        "    )\n",
        "    return True"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gyJEPeBlIHXf",
        "colab_type": "code",
        "outputId": "6973c212-7886-488d-a872-09949e67cdcc",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "compare_fwd(\n",
        "    reset_dummy_init(model_v1()),\n",
        "    reset_dummy_init(model_v2b()),\n",
        ")"
      ],
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 20
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fqE5GfH5M_vK",
        "colab_type": "text"
      },
      "source": [
        "This is promising. It remains to check that the initialisations of the two models agree. In general this can be somewhat tricky for the reasons given above. The current situation is actually much nicer as both models are initialised with a final call to an `init_cnn` function and if we fix the random seed before this call we get agreement:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tXK2ZhWZWwgK",
        "colab_type": "code",
        "outputId": "f0e1e08d-0aed-4d93-ed7a-5cc1d484b075",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "model1 = model_v1()\n",
        "torch.manual_seed(123)\n",
        "mxresnet.init_cnn(model1)\n",
        "\n",
        "model2 = model_v2b()\n",
        "torch.manual_seed(123)\n",
        "fastai2.vision.models.xresnet.init_cnn(model2)\n",
        "\n",
        "compare_fwd(model1, model2)"
      ],
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 21
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ivWSZvNSbGK8",
        "colab_type": "text"
      },
      "source": [
        "### Fastai v2 training"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yi7bhGeBm5mc",
        "colab_type": "text"
      },
      "source": [
        "Here is a first attempt at training using the v2 model and codebase + the DALI dataloader we used above. We will use the ranger optimiser from the v1 codebase for now:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mPi75pwkftVr",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import fastai2.callback.all #need to import this to patch fastai2.basics.Learner with .to_fp16() method\n",
        "RangerWrapper = lambda *args, **kwargs: fastai2.basics.OptimWrapper(ranger.Ranger(*args, **kwargs))\n",
        "\n",
        "dali_data_v2 = lambda data_dir=data_dir, bs=bs: fastai2.basics.DataBunch(train_dl(data_dir/'train', bs), valid_dl(data_dir/'val', bs))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "K2WchF9Qwcm-",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model_v2 = partial(xresnet18, c_out=10, sa=1, sym=0, act_cls=(lambda: mish))\n",
        "\n",
        "learner_v2 = partial(fastai2.basics.Learner, \n",
        "                  lr=4e-3,\n",
        "                  opt_func=partial(RangerWrapper, betas=(0.95, 0.99), eps=1e-6),\n",
        "                  metrics=(fastai2.metrics.accuracy,),\n",
        "                  loss_func=fastai2.basics.LabelSmoothingCrossEntropy(),\n",
        "                  )"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0rMgOupgFSJa",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 204
        },
        "outputId": "30aa23b2-6014-489e-d3a2-a2d84a94eeca"
      },
      "source": [
        "learn = learner_v2(dali_data_v2(), model_v2()).to_fp16().fit_flat_cos(5, pct_start=0.72)"
      ],
      "execution_count": 24,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: left;\">\n",
              "      <th>epoch</th>\n",
              "      <th>train_loss</th>\n",
              "      <th>valid_loss</th>\n",
              "      <th>accuracy</th>\n",
              "      <th>time</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <td>0</td>\n",
              "      <td>1.442993</td>\n",
              "      <td>1.237608</td>\n",
              "      <td>0.697000</td>\n",
              "      <td>00:19</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>1</td>\n",
              "      <td>1.199762</td>\n",
              "      <td>1.103699</td>\n",
              "      <td>0.767000</td>\n",
              "      <td>00:19</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>2</td>\n",
              "      <td>1.110460</td>\n",
              "      <td>1.061970</td>\n",
              "      <td>0.785000</td>\n",
              "      <td>00:19</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>3</td>\n",
              "      <td>1.019471</td>\n",
              "      <td>0.928476</td>\n",
              "      <td>0.859000</td>\n",
              "      <td>00:19</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>4</td>\n",
              "      <td>0.916012</td>\n",
              "      <td>0.843927</td>\n",
              "      <td>0.883000</td>\n",
              "      <td>00:19</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dQBHVophoeAv",
        "colab_type": "text"
      },
      "source": [
        "To be continued..."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F7CYZ6dIRRzy",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}