conversion.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "convmixer-conversion.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyNsckeGOnqgONpT3IdP95Md",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Rishit-dagli/ConvMixer-torch2tf/blob/main/conversion.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WABtusco3WjU"
      },
      "source": [
        "# Conversion Scripts\n",
        "\n",
        "Authors: [Rishit Dagli](https://twitter.com/rishit_dagli)\n",
        "\n",
        "In this Notebook we convert these 3 ConvMixer models to TensorFlow SavedModels:\n",
        "\n",
        "| Model Name | resolution | acc@1 | #params | File Size | TensorFlow Model |\n",
        "|------------|:---:|:---:|:---:|----------:|:--------------:|\n",
        "| ConvMixer-1536/20 | 224x224 | 81.37 | 51.6 x 10^6 | 184MB | [github](https://github.com/Rishit-dagli/ConvMixer-torch2tf/releases/download/untagged-de81892a1b06347b8d97/convmixer_1536_20.tar.gz)/[drive](https://drive.google.com/file/d/1qrzap4vi2KFQTHxf9h_AMbWGvtbP5rIA/view?usp=sharing)/[bucket](https://storage.googleapis.com/convmixer-hubmodels.appspot.com/convmixer_1536_20.tar.gz) |\n",
        "| ConvMixer-768/32 | 224x224 | 80.16 | 21.1 x 10^6 | 75MB | [github](https://github.com/Rishit-dagli/ConvMixer-torch2tf/releases/download/untagged-de81892a1b06347b8d97/convmixer_768_32.tar.gz)/[drive](https://drive.google.com/file/d/1NJgHKjPd3YC8XHypQIs5A05XKd15o0s3/view?usp=sharing)/[bucket](https://storage.googleapis.com/convmixer-hubmodels.appspot.com/convmixer_768_32.tar.gz) |\n",
        "| ConvMixer-1024/20 | 224x224 | 76.94 | 24.4 x 10^6 | 87MB | [github](https://github.com/Rishit-dagli/ConvMixer-torch2tf/releases/download/untagged-de81892a1b06347b8d97/convmixer_1024_20.tar.gz)/[drive](https://drive.google.com/file/d/1--jRgK0KmLtWCJswYtfxSIfEcjAOrJyv/view?usp=sharing)/[bucket](https://storage.googleapis.com/convmixer-hubmodels.appspot.com/convmixer_1024_20.tar.gz) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bbzxwe_EyHhF"
      },
      "source": [
        "## Setup\n",
        "\n",
        "We start off by installing [`timm`](https://github.com/rwightman/pytorch-image-models) but notice we are installing from my fork of it, this fork includes code from `timm_convmixer.py` from this repo. This updates the way the model handles padding in a `nn.Conv2D` layer to allow it to be converted."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "T946qhB9wjJe",
        "outputId": "a5a40433-c093-4477-c2d6-f2b24ffb91bf"
      },
      "source": [
        "!pip install git+https://github.com/Rishit-dagli/pytorch-image-models.git@Rishit-dagli-convmixer-tf#egg=timm"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting timm\n",
            "  Cloning https://github.com/Rishit-dagli/pytorch-image-models.git (to revision Rishit-dagli-convmixer-tf) to /tmp/pip-install-e8uqgzff/timm_e597bf84cf2d47059e0ea3746c4ec313\n",
            "  Running command git clone -q https://github.com/Rishit-dagli/pytorch-image-models.git /tmp/pip-install-e8uqgzff/timm_e597bf84cf2d47059e0ea3746c4ec313\n",
            "  Running command git checkout -b Rishit-dagli-convmixer-tf --track origin/Rishit-dagli-convmixer-tf\n",
            "  Switched to a new branch 'Rishit-dagli-convmixer-tf'\n",
            "  Branch 'Rishit-dagli-convmixer-tf' set up to track remote branch 'Rishit-dagli-convmixer-tf' from 'origin'.\n",
            "Requirement already satisfied: torch>=1.4 in /usr/local/lib/python3.7/dist-packages (from timm) (1.9.0+cu111)\n",
            "Requirement already satisfied: torchvision in /usr/local/lib/python3.7/dist-packages (from timm) (0.10.0+cu111)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch>=1.4->timm) (3.7.4.3)\n",
            "Requirement already satisfied: pillow>=5.3.0 in /usr/local/lib/python3.7/dist-packages (from torchvision->timm) (7.1.2)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from torchvision->timm) (1.19.5)\n",
            "Building wheels for collected packages: timm\n",
            "  Building wheel for timm (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for timm: filename=timm-0.5.0-py3-none-any.whl size=417714 sha256=356aa1cab8ace2c32f2536b6851573a6721558de533feffb56e0630625d508b9\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-astx_hvu/wheels/29/a9/a4/4ad1e14f936318e78dec16e4541708836546145fafa44e5915\n",
            "Successfully built timm\n",
            "Installing collected packages: timm\n",
            "Successfully installed timm-0.5.0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2zf68VqOzBbW",
        "outputId": "297c5f3b-9f17-4402-b402-8b6da14bface"
      },
      "source": [
        "!pip install --pre torch torchvision torchaudio -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html --upgrade"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in links: https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html\n",
            "Requirement already satisfied: torch in /usr/local/lib/python3.7/dist-packages (1.9.0+cu111)\n",
            "Collecting torch\n",
            "  Downloading https://download.pytorch.org/whl/nightly/cpu/torch-1.11.0.dev20211023%2Bcpu-cp37-cp37m-linux_x86_64.whl (151.2 MB)\n",
            "\u001b[K     |████████████████████████████████| 151.2 MB 17 kB/s \n",
            "\u001b[?25hRequirement already satisfied: torchvision in /usr/local/lib/python3.7/dist-packages (0.10.0+cu111)\n",
            "Collecting torchvision\n",
            "  Downloading https://download.pytorch.org/whl/nightly/cpu/torchvision-0.12.0.dev20211023%2Bcpu-cp37-cp37m-linux_x86_64.whl (16.8 MB)\n",
            "\u001b[K     |████████████████████████████████| 16.8 MB 59 kB/s \n",
            "\u001b[?25hCollecting torchaudio\n",
            "  Downloading https://download.pytorch.org/whl/nightly/cpu/torchaudio-0.11.0.dev20211023%2Bcpu-cp37-cp37m-linux_x86_64.whl (2.7 MB)\n",
            "\u001b[K     |████████████████████████████████| 2.7 MB 44.8 MB/s \n",
            "\u001b[?25hRequirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch) (3.7.4.3)\n",
            "Requirement already satisfied: pillow!=8.3.0,>=5.3.0 in /usr/local/lib/python3.7/dist-packages (from torchvision) (7.1.2)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from torchvision) (1.19.5)\n",
            "Installing collected packages: torch, torchvision, torchaudio\n",
            "  Attempting uninstall: torch\n",
            "    Found existing installation: torch 1.9.0+cu111\n",
            "    Uninstalling torch-1.9.0+cu111:\n",
            "      Successfully uninstalled torch-1.9.0+cu111\n",
            "  Attempting uninstall: torchvision\n",
            "    Found existing installation: torchvision 0.10.0+cu111\n",
            "    Uninstalling torchvision-0.10.0+cu111:\n",
            "      Successfully uninstalled torchvision-0.10.0+cu111\n",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "torchtext 0.10.0 requires torch==1.9.0, but you have torch 1.11.0.dev20211023+cpu which is incompatible.\u001b[0m\n",
            "Successfully installed torch-1.11.0.dev20211023+cpu torchaudio-0.11.0.dev20211023+cpu torchvision-0.12.0.dev20211023+cpu\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QsAJxZqg1Ra4",
        "outputId": "5c0f71a9-38df-44c6-e28e-c3026dde0e4c"
      },
      "source": [
        "!pip install onnx onnx-tf"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting onnx\n",
            "  Downloading onnx-1.10.1-cp37-cp37m-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (12.3 MB)\n",
            "\u001b[K     |████████████████████████████████| 12.3 MB 7.3 MB/s \n",
            "\u001b[?25hCollecting onnx-tf\n",
            "  Downloading onnx_tf-1.9.0-py3-none-any.whl (222 kB)\n",
            "\u001b[K     |████████████████████████████████| 222 kB 64.5 MB/s \n",
            "\u001b[?25hRequirement already satisfied: numpy>=1.16.6 in /usr/local/lib/python3.7/dist-packages (from onnx) (1.19.5)\n",
            "Requirement already satisfied: protobuf in /usr/local/lib/python3.7/dist-packages (from onnx) (3.17.3)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from onnx) (1.15.0)\n",
            "Requirement already satisfied: typing-extensions>=3.6.2.1 in /usr/local/lib/python3.7/dist-packages (from onnx) (3.7.4.3)\n",
            "Collecting tensorflow-addons\n",
            "  Downloading tensorflow_addons-0.14.0-cp37-cp37m-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (1.1 MB)\n",
            "\u001b[K     |████████████████████████████████| 1.1 MB 43.3 MB/s \n",
            "\u001b[?25hRequirement already satisfied: PyYAML in /usr/local/lib/python3.7/dist-packages (from onnx-tf) (3.13)\n",
            "Requirement already satisfied: typeguard>=2.7 in /usr/local/lib/python3.7/dist-packages (from tensorflow-addons->onnx-tf) (2.7.1)\n",
            "Installing collected packages: tensorflow-addons, onnx, onnx-tf\n",
            "Successfully installed onnx-1.10.1 onnx-tf-1.9.0 tensorflow-addons-0.14.0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wtx6V6Zz5jDN"
      },
      "source": [
        "import os\n",
        "from io import BytesIO\n",
        "\n",
        "import numpy as np\n",
        "import onnx\n",
        "import requests\n",
        "import tensorflow as tf\n",
        "import tensorflow_hub as hub\n",
        "import timm\n",
        "import torch\n",
        "from onnx_tf.backend import prepare\n",
        "from PIL import Image"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9nqO5Bp6A7h-"
      },
      "source": [
        "## Choose pretrained model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qAIg0mTSzmgI"
      },
      "source": [
        "# fmt: off\n",
        "\n",
        "#@title Choose a model type\n",
        "model_variant = \"convmixer_1024_20_ks9_p14\" #@param ['convmixer_1536_20', 'convmixer_768_32', 'convmixer_1024_20_ks9_p14']\n",
        "resolution = [224, 224] \n",
        "num_classes = 1000\n",
        "\n",
        "# fmt: on"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cbBQ209bA-SH"
      },
      "source": [
        "## Loading PyTorch pre-trained Model\n",
        "\n",
        "We will now load the pre-trained PyTorch models."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kLuWhSGKkyzK",
        "outputId": "71ffc4a4-d9c4-466d-e099-94a770a2f813"
      },
      "source": [
        "model = timm.create_model(model_variant, pretrained=True)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Downloading: \"https://github.com/tmp-iclr/convmixer/releases/download/timm-v1.0/convmixer_1024_20_ks9_p14.pth.tar\" to /root/.cache/torch/hub/checkpoints/convmixer_1024_20_ks9_p14.pth.tar\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CFwVkfn78bg1"
      },
      "source": [
        "### Transpose channel axes\n",
        "\n",
        "Often in PyTorch we use the channels first data format meaning you would want to pass in input $(N, C_{in}, H_{in}, W_{in})$ wwhereas in TensorFlow you would often want to use the channels last format; pass in input as $(N, H_{in}, W_{in}, C_{in})$. You can also notice that the default data format for [nn.Conv2d](https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html#torch.nn.Conv2d) is channels first and for [tf.keras.layers.Conv2D](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Conv2D) is channels last.\n",
        "\n",
        "To do so we create a wrapper model over the earlier model we created which just transposes the channel axes in the input tensor."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UeSjQjYyR-Sc"
      },
      "source": [
        "class TransposeModelWrapper(torch.nn.Module):\n",
        "    def __init__(self, main_module):\n",
        "        super(TransposeModelWrapper, self).__init__()\n",
        "        self.main_module = main_module\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = torch.transpose(x, 1, 3)\n",
        "        return self.main_module(x)\n",
        "\n",
        "\n",
        "transposed_model = TransposeModelWrapper(model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8IiB5zPMA3U8"
      },
      "source": [
        "## Convert to ONNX\n",
        "\n",
        "We set `dynamic_axes` to allow the converted model operate with arbitrary batch sizes."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OkK6os42l3W6"
      },
      "source": [
        "batch_size = 1\n",
        "x = torch.randn(batch_size, resolution[0], resolution[1], 3)\n",
        "model_onnx = model_variant + \".onnx\"\n",
        "\n",
        "torch.onnx.export(\n",
        "    transposed_model,\n",
        "    x,\n",
        "    model_onnx,\n",
        "    opset_version=9,\n",
        "    input_names=[\"input\"],\n",
        "    output_names=[\"output\"],\n",
        "    dynamic_axes={\"input\": {0: \"batch_size\"}, \"output\": {0: \"batch_size\"}},\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kLSGDi0PBCee"
      },
      "source": [
        "## Create a TensorFlow SavedModel\n",
        "\n",
        "We will first load the onnx model and then check that the IR is well formed."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NvER3xjP0w7L"
      },
      "source": [
        "onnx_model = onnx.load(model_onnx)\n",
        "onnx.checker.check_model(onnx_model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hoomhUy8-uji"
      },
      "source": [
        "We will now use the ONNX TensorFlow backend to convert our model to a TensorFlow SavedModel."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hsD-0BMX4P8e"
      },
      "source": [
        "os.environ['onnx_model_dir'] = \"/content/\" + model_variant + \".onnx\"\n",
        "os.environ['saved_model_dir'] = \"/content/tf_\" + model_variant"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "V3SDlDUuqDlf",
        "outputId": "dbdda1a9-4a5d-4337-e345-5aca6f63ca8c"
      },
      "source": [
        "!onnx-tf convert -i $onnx_model_dir -o $saved_model_dir"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "2021-10-23 16:06:25,111 - onnx-tf - INFO - Start converting onnx pb to tf saved model\n",
            "2021-10-23 16:06:25.553011: E tensorflow/stream_executor/cuda/cuda_driver.cc:271] failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected\n",
            "2021-10-23 16:06:25.553097: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:156] kernel driver does not appear to be running on this host (b225fa41d07c): /proc/driver/nvidia/version does not exist\n",
            "2021-10-23 16:06:38.647195: W tensorflow/python/util/util.cc:348] Sets are not currently considered sequences, but this may change in the future, so consider avoiding using them.\n",
            "WARNING:absl:Found untraced functions such as gen_tensor_dict while saving (showing 1 of 1). These functions will not be directly callable after loading.\n",
            "2021-10-23 16:15:32,177 - onnx-tf - INFO - Converting completes successfully.\n",
            "INFO:onnx-tf:Converting completes successfully.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BNL-EIFTBGK1"
      },
      "source": [
        "## Verify SavedModel works as expected"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5-Ns6QFJ4-Qm"
      },
      "source": [
        "saved_model_dir = \"/content/tf_\" + model_variant"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e_ETwInCBI5Q"
      },
      "source": [
        "### Load SavedModel"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xkH4gI6yBYvE"
      },
      "source": [
        "model = tf.saved_model.load(saved_model_dir)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "j5HqqOiQAI66"
      },
      "source": [
        "layer = hub.KerasLayer(\n",
        "    saved_model_dir, signature=\"serving_default\", signature_outputs_as_dict=True\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Uj9K137JD4QL"
      },
      "source": [
        "### Check input TypeSpec\n",
        "\n",
        "We will now use the SavedModel CLI to verify that our SavedModel has the correct input type spec which should be $(-1, 224, 224, 3)$ depending on which model you converted."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AxrJpELSBz1_",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "f6f09e30-948c-43c2-c1bd-f9b684a55340"
      },
      "source": [
        "!saved_model_cli show --dir $saved_model_dir --tag_set serve --signature_def serving_default"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "The given SavedModel SignatureDef contains the following input(s):\n",
            "  inputs['input'] tensor_info:\n",
            "      dtype: DT_FLOAT\n",
            "      shape: (-1, 224, 224, 3)\n",
            "      name: serving_default_input:0\n",
            "The given SavedModel SignatureDef contains the following output(s):\n",
            "  outputs['output'] tensor_info:\n",
            "      dtype: DT_FLOAT\n",
            "      shape: (-1, 1000)\n",
            "      name: PartitionedCall:0\n",
            "Method name is: tensorflow/serving/predict\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-Yh4fgJqKib-"
      },
      "source": [
        "## Inference with the model\n",
        "\n",
        "We will now try inferencing with our TensorFlow SavedModel first on a dummy tensor and then on an image.\n",
        "\n",
        "### Dummy Tensor\n",
        "\n",
        "In the first example we try on a dummy tensor with $N=5$ to verify using an arbitary batch size works well. The output of these code cells should be:\n",
        "\n",
        "- `TensorShape([5, 1000])`\n",
        "- `TensorShape([1, 1000])`\n",
        "\n",
        "depending on which model you convert."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fm4APC-XOSPJ",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "7d670ce2-d71d-4e51-a28b-59450b889ce2"
      },
      "source": [
        "dummy = model.signatures[\"serving_default\"](\n",
        "    tf.random.normal([5, resolution[0], resolution[1], 3])\n",
        ")\n",
        "list(dummy.values())[0].shape"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "TensorShape([5, 1000])"
            ]
          },
          "metadata": {},
          "execution_count": 16
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WAtqFAl5SCp2",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "9d701da8-becb-4baf-9232-f395461446ff"
      },
      "source": [
        "dummy = layer(tf.random.normal([1, resolution[0], resolution[1], 3]))\n",
        "list(dummy.values())[0].shape"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "TensorShape([1, 1000])"
            ]
          },
          "metadata": {},
          "execution_count": 17
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LKknMBe9Z9im"
      },
      "source": [
        "### On a real image\n",
        "\n",
        "#### Image preprocessing utilities (adapted from [Willi Gierke](https://ch.linkedin.com/in/willi-gierke))"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gSv8fL4yK5xo",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "7934f58d-0d11-4aa9-8fb1-2c33f1cf8a84"
      },
      "source": [
        "def preprocess_image(image):\n",
        "    image = np.array(image)\n",
        "    image_resized = tf.image.resize(image, (resolution[0], resolution[1]))\n",
        "    image_resized = tf.cast(image_resized, tf.float32)\n",
        "    image_resized = image_resized / 255\n",
        "    image_resized = tf.keras.layers.Normalization(\n",
        "        mean=(0.485, 0.456, 0.406), variance=(0.052441, 0.050176, 0.050625)\n",
        "    )(image_resized)\n",
        "    return tf.expand_dims(image_resized, 0).numpy()\n",
        "\n",
        "\n",
        "def load_image_from_url(url):\n",
        "    response = requests.get(url)\n",
        "    image = Image.open(BytesIO(response.content))\n",
        "    image = preprocess_image(image)\n",
        "    return image\n",
        "    \n",
        "!wget https://storage.googleapis.com/bit_models/ilsvrc2012_wordnet_lemmas.txt -O ilsvrc2012_wordnet_lemmas.txt"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--2021-10-23 16:15:53--  https://storage.googleapis.com/bit_models/ilsvrc2012_wordnet_lemmas.txt\n",
            "Resolving storage.googleapis.com (storage.googleapis.com)... 74.125.142.128, 74.125.195.128, 173.194.202.128, ...\n",
            "Connecting to storage.googleapis.com (storage.googleapis.com)|74.125.142.128|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 21675 (21K) [text/plain]\n",
            "Saving to: ‘ilsvrc2012_wordnet_lemmas.txt’\n",
            "\n",
            "\r          ilsvrc201   0%[                    ]       0  --.-KB/s               \rilsvrc2012_wordnet_ 100%[===================>]  21.17K  --.-KB/s    in 0s      \n",
            "\n",
            "2021-10-23 16:15:53 (73.5 MB/s) - ‘ilsvrc2012_wordnet_lemmas.txt’ saved [21675/21675]\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Pb0qyE5SBQwc"
      },
      "source": [
        "#### Load image and infer"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CvrZzFyMNo2j"
      },
      "source": [
        "with open(\"ilsvrc2012_wordnet_lemmas.txt\", \"r\") as f:\n",
        "    lines = f.readlines()\n",
        "imagenet_int_to_str = [line.rstrip() for line in lines]\n",
        "\n",
        "def infer_on_image(img_url, expected_label):\n",
        "    image = load_image_from_url(img_url)\n",
        "    predictions = model.signatures[\"serving_default\"](tf.constant(image))\n",
        "    logits = predictions[\"output\"][0]\n",
        "    predicted_label = imagenet_int_to_str[int(np.argmax(logits))]\n",
        "    assert (\n",
        "        predicted_label == expected_label\n",
        "    ), f\"Expected {expected_label} but was {predicted_label}\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cBB5UfU7xWBW"
      },
      "source": [
        "infer_on_image(img_url = \"https://i.imgur.com/IKR2LyP.jpg\", expected_label = \"purse\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lRXCGZN79VKV"
      },
      "source": [
        "As a sanity test lets try 5 more images."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "egTvfvrE_MWt"
      },
      "source": [
        "infer_on_image(img_url = \"https://i.imgur.com/mH0Wrvb.jpg\", expected_label = \"goldfish, Carassius_auratus\")\n",
        "infer_on_image(img_url = \"https://i.imgur.com/A5m4ZG1.jpg\", expected_label = \"scorpion\")\n",
        "infer_on_image(img_url = \"https://i.imgur.com/faOAEFg.jpg\", expected_label = \"leatherback_turtle, leatherback, leathery_turtle, Dermochelys_coriacea\")\n",
        "infer_on_image(img_url = \"https://i.imgur.com/lfhdaSi.jpg\", expected_label = \"Siamese_cat, Siamese\")\n",
        "infer_on_image(img_url = \"https://i.imgur.com/Qwa8wHX.jpg\", expected_label = \"boa_constrictor, Constrictor_constrictor\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hu7WGw-fiBAg"
      },
      "source": [
        "## References\n",
        "\n",
        "[1] Anonymous. Patches Are All You Need? 2021. openreview.net, https://openreview.net/forum?id=TVHS5Y4dNvM.\n",
        "\n",
        "[2] Official Code Implementation: https://github.com/tmp-iclr/convmixer\n",
        "\n",
        "[3] Ross Wightman, . \"PyTorch Image Models.\" https://github.com/rwightman/pytorch-image-models."
      ]
    }
  ]
}