pytorch/pytorch_to_all_pipeline.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Notes**:\n",
    "\n",
    "Success except for tflite"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys, os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'3.6.5 |Anaconda, Inc.| (default, Apr 29 2018, 16:14:56) \\n[GCC 7.2.0]'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sys.version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastai.vision import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1.0.1.post2'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'0.2.1'"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torchvision.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1.13.1'"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import onnx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1.4.1'"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "onnx.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['BasicBlock',\n",
       " 'Darknet',\n",
       " 'DynamicUnet',\n",
       " 'ResLayer',\n",
       " 'ResNet',\n",
       " 'SqueezeNet',\n",
       " 'UnetBlock',\n",
       " 'WideResNet',\n",
       " 'XResNet',\n",
       " 'alexnet',\n",
       " 'darknet',\n",
       " 'densenet121',\n",
       " 'densenet161',\n",
       " 'densenet169',\n",
       " 'densenet201',\n",
       " 'resnet101',\n",
       " 'resnet152',\n",
       " 'resnet18',\n",
       " 'resnet34',\n",
       " 'resnet50',\n",
       " 'squeezenet1_0',\n",
       " 'squeezenet1_1',\n",
       " 'unet',\n",
       " 'vgg16_bn',\n",
       " 'vgg19_bn',\n",
       " 'wrn',\n",
       " 'wrn_22',\n",
       " 'xception',\n",
       " 'xresnet',\n",
       " 'xresnet101',\n",
       " 'xresnet152',\n",
       " 'xresnet18',\n",
       " 'xresnet34',\n",
       " 'xresnet50']"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[m for m in dir(models) if m[:2] != '__']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.models import vgg11\n",
    "models.vgg11 = vgg11"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODELS = Path('./models/'); MODELS.ls()\n",
    "arch_name = 'vgg11'\n",
    "\n",
    "def inputs_same(x1,x2): return np.all(np.isclose(x1, x2, atol=1e-5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading: \"https://download.pytorch.org/models/vgg11-bbd30ac9.pth\" to /home/ubuntu/.torch/models/vgg11-bbd30ac9.pth\n",
      "531456000it [00:33, 15797853.37it/s]\n"
     ]
    }
   ],
   "source": [
    "f = getattr(models, arch_name)\n",
    "model = f(True); model.eval();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Add permute for tflite"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "class Model(nn.Module):\n",
    "    def __init__(self, model=model):\n",
    "        super().__init__()\n",
    "        self.model = model\n",
    "    \n",
    "    def forward(self, x):    \n",
    "        x = x.permute(0,3,1,2).contiguous()\n",
    "        return self.model(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Model(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Model(\n",
       "  (model): VGG(\n",
       "    (features): Sequential(\n",
       "      (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (1): ReLU(inplace)\n",
       "      (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "      (3): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (4): ReLU(inplace)\n",
       "      (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "      (6): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (7): ReLU(inplace)\n",
       "      (8): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (9): ReLU(inplace)\n",
       "      (10): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "      (11): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (12): ReLU(inplace)\n",
       "      (13): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (14): ReLU(inplace)\n",
       "      (15): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "      (16): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (17): ReLU(inplace)\n",
       "      (18): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      (19): ReLU(inplace)\n",
       "      (20): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "    )\n",
       "    (classifier): Sequential(\n",
       "      (0): Linear(in_features=25088, out_features=4096, bias=True)\n",
       "      (1): ReLU(inplace)\n",
       "      (2): Dropout(p=0.5)\n",
       "      (3): Linear(in_features=4096, out_features=4096, bias=True)\n",
       "      (4): ReLU(inplace)\n",
       "      (5): Dropout(p=0.5)\n",
       "      (6): Linear(in_features=4096, out_features=1000, bias=True)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 96,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastai.layers import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 156,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=0)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = x.permute(0,3,1,2).contiguous()\n",
    "        return self.conv1(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 157,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = SimpleModel()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get input ready"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 158,
   "metadata": {},
   "outputs": [],
   "source": [
    "img = open_image(\"cat224x224.jpg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 159,
   "metadata": {},
   "outputs": [],
   "source": [
    "normalized_img = normalize(img.data, tensor(imagenet_stats[0]), tensor(imagenet_stats[1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 160,
   "metadata": {},
   "outputs": [],
   "source": [
    "# batch x channel x height x width\n",
    "torch_input = normalized_img[None, ...]\n",
    "numpy_input = to_np(torch_input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 161,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 3, 224, 224])"
      ]
     },
     "execution_count": 161,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch_input.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 162,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch_input = torch_input.permute(0,2,3,1)\n",
    "numpy_input = to_np(torch_input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 163,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([1, 224, 224, 3]), (1, 224, 224, 3))"
      ]
     },
     "execution_count": 163,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch_input.shape, numpy_input.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### pytorch output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 164,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "    with torch.cuda.device(0):\n",
    "        torch_output = to_np(model(torch_input))\n",
    "else: torch_output = to_np(model(torch_input))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 165,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2063640"
      ]
     },
     "execution_count": 165,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.argmax(torch_output[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 166,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[ 5.928250e-01,  6.495560e-01,  6.586254e-01,  6.328704e-01, ...,  1.145797e-01, -1.446192e-02,\n",
       "         -9.115485e-02, -1.700499e-01],\n",
       "        [ 6.237740e-01,  6.704966e-01,  6.704947e-01,  6.700016e-01, ...,  1.083156e-01, -3.941596e-02,\n",
       "         -9.333248e-02, -1.599272e-01],\n",
       "        [ 6.543528e-01,  7.007562e-01,  6.936352e-01,  7.207102e-01, ...,  1.495145e-01, -5.462829e-04,\n",
       "         -5.724194e-02, -1.157975e-01],\n",
       "        [ 6.586937e-01,  6.734817e-01,  6.918388e-01,  7.479682e-01, ...,  1.761432e-01,  5.087249e-02,\n",
       "         -1.233188e-02, -3.808299e-02],\n",
       "        ...,\n",
       "        [-5.286742e-01, -3.994761e-01, -4.185633e-01, -5.430715e-01, ..., -8.238600e-01, -8.971226e-01,\n",
       "         -1.071087e+00, -1.090646e+00],\n",
       "        [-6.421031e-01, -5.654278e-01, -4.625924e-01, -5.476409e-01, ..., -8.586741e-01, -7.533782e-01,\n",
       "         -8.619232e-01, -9.086895e-01],\n",
       "        [-6.917681e-01, -6.695323e-01, -5.033025e-01, -4.304609e-01, ..., -8.144889e-01, -6.926728e-01,\n",
       "         -7.351253e-01, -7.078845e-01],\n",
       "        [-5.975550e-01, -6.254424e-01, -5.287780e-01, -4.007401e-01, ..., -7.700869e-01, -7.775310e-01,\n",
       "         -9.107289e-01, -8.765057e-01]],\n",
       "\n",
       "       [[ 9.304108e-01,  9.061006e-01,  8.904200e-01,  8.748452e-01, ...,  1.507193e-01,  4.863602e-02,\n",
       "         -2.965286e-02, -6.887357e-02],\n",
       "        [ 9.354134e-01,  9.211414e-01,  9.025218e-01,  8.819546e-01, ...,  1.583953e-01,  7.783134e-02,\n",
       "          2.384748e-03, -2.762413e-02],\n",
       "        [ 9.247492e-01,  9.207308e-01,  9.197693e-01,  9.268373e-01, ...,  1.964801e-01,  1.070963e-01,\n",
       "          6.585616e-02,  3.885203e-02],\n",
       "        [ 8.749162e-01,  9.000026e-01,  9.111722e-01,  9.365388e-01, ...,  2.133995e-01,  1.541420e-01,\n",
       "          1.271267e-01,  1.043288e-01],\n",
       "        ...,\n",
       "        [-4.185115e-01, -4.692028e-01, -4.844882e-01, -4.011375e-01, ..., -9.281288e-01, -9.043365e-01,\n",
       "         -8.730649e-01, -8.296068e-01],\n",
       "        [-4.699255e-01, -4.824458e-01, -4.375248e-01, -3.610857e-01, ..., -6.990985e-01, -7.424051e-01,\n",
       "         -7.382967e-01, -6.652063e-01],\n",
       "        [-5.109918e-01, -4.493734e-01, -4.005437e-01, -4.306577e-01, ..., -6.778324e-01, -7.774227e-01,\n",
       "         -8.163334e-01, -7.388382e-01],\n",
       "        [-4.955843e-01, -4.066808e-01, -4.318877e-01, -5.693138e-01, ..., -6.130958e-01, -7.152777e-01,\n",
       "         -7.512010e-01, -6.329308e-01]],\n",
       "\n",
       "       [[-8.835173e-01, -8.811719e-01, -8.571362e-01, -8.381184e-01, ..., -1.226337e-01, -1.413189e-02,\n",
       "          6.252305e-02,  1.186951e-01],\n",
       "        [-8.997300e-01, -9.013266e-01, -8.836781e-01, -8.679811e-01, ..., -1.311588e-01, -3.768335e-02,\n",
       "          2.365048e-02,  6.502268e-02],\n",
       "        [-8.955600e-01, -9.047340e-01, -9.121482e-01, -9.156053e-01, ..., -1.767565e-01, -9.044454e-02,\n",
       "         -3.904480e-02, -1.396201e-03],\n",
       "        [-8.667954e-01, -8.885559e-01, -9.212384e-01, -9.478627e-01, ..., -2.214962e-01, -1.539418e-01,\n",
       "         -1.061586e-01, -7.190461e-02],\n",
       "        ...,\n",
       "        [ 6.527845e-01,  6.506559e-01,  6.483269e-01,  6.090642e-01, ...,  1.127803e+00,  1.157625e+00,\n",
       "          1.176368e+00,  1.130930e+00],\n",
       "        [ 7.203699e-01,  6.666610e-01,  6.309564e-01,  6.287088e-01, ...,  9.141809e-01,  9.502071e-01,\n",
       "          9.911880e-01,  9.432814e-01],\n",
       "        [ 7.402143e-01,  6.619299e-01,  6.260362e-01,  6.843429e-01, ...,  9.217322e-01,  9.615628e-01,\n",
       "          1.000422e+00,  9.200149e-01],\n",
       "        [ 7.193744e-01,  6.757457e-01,  6.706685e-01,  7.347934e-01, ...,  9.194486e-01,  9.976853e-01,\n",
       "          1.031135e+00,  9.027116e-01]],\n",
       "\n",
       "       [[-5.611107e-01, -6.185374e-01, -6.274537e-01, -6.018535e-01, ..., -3.286290e-01, -1.958319e-01,\n",
       "         -1.233654e-01, -5.534839e-02],\n",
       "        [-5.882207e-01, -6.269063e-01, -6.223419e-01, -5.886477e-01, ..., -2.825735e-01, -1.406765e-01,\n",
       "         -6.240558e-02, -1.910090e-03],\n",
       "        [-6.440452e-01, -6.522714e-01, -6.412833e-01, -5.956016e-01, ..., -2.842835e-01, -1.286600e-01,\n",
       "         -5.839922e-02, -1.226380e-02],\n",
       "        [-6.646396e-01, -6.643090e-01, -6.534098e-01, -6.352038e-01, ..., -3.086022e-01, -1.643063e-01,\n",
       "         -1.076260e-01, -7.506974e-02],\n",
       "        ...,\n",
       "        [ 1.974958e-01,  1.334879e-01,  2.476262e-01,  3.490084e-01, ...,  9.260128e-01,  9.112363e-01,\n",
       "          8.825827e-01,  8.980755e-01],\n",
       "        [ 2.174751e-01,  2.777752e-01,  3.892178e-01,  3.058348e-01, ...,  6.894872e-01,  6.306827e-01,\n",
       "          6.324791e-01,  7.319039e-01],\n",
       "        [ 4.317094e-01,  5.105460e-01,  4.332957e-01,  2.208417e-01, ...,  6.646112e-01,  5.065519e-01,\n",
       "          4.617480e-01,  6.012140e-01],\n",
       "        [ 4.610324e-01,  4.142273e-01,  2.067479e-01,  1.124400e-01, ...,  3.580654e-01,  2.135994e-01,\n",
       "          1.807216e-01,  3.309866e-01]],\n",
       "\n",
       "       ...,\n",
       "\n",
       "       [[-1.021994e+00, -1.007560e+00, -1.020993e+00, -1.028300e+00, ..., -6.495625e-01, -6.656052e-01,\n",
       "         -6.676931e-01, -6.754252e-01],\n",
       "        [-1.022150e+00, -1.020793e+00, -1.031548e+00, -1.039426e+00, ..., -6.874483e-01, -6.935481e-01,\n",
       "         -7.095402e-01, -7.143820e-01],\n",
       "        [-1.000767e+00, -1.022446e+00, -1.045676e+00, -1.072399e+00, ..., -7.137644e-01, -7.332122e-01,\n",
       "         -7.495062e-01, -7.635692e-01],\n",
       "        [-9.922659e-01, -1.020445e+00, -1.055791e+00, -1.083923e+00, ..., -7.817140e-01, -7.956998e-01,\n",
       "         -8.141065e-01, -8.214797e-01],\n",
       "        ...,\n",
       "        [ 5.825371e-01,  5.784728e-01,  5.523584e-01,  5.299688e-01, ...,  1.002575e+00,  1.063073e+00,\n",
       "          1.038671e+00,  9.423774e-01],\n",
       "        [ 5.757274e-01,  5.073735e-01,  4.978513e-01,  5.697755e-01, ...,  7.133898e-01,  8.178766e-01,\n",
       "          8.427096e-01,  7.569195e-01],\n",
       "        [ 6.112009e-01,  5.420885e-01,  5.843502e-01,  6.813412e-01, ...,  8.724557e-01,  9.414760e-01,\n",
       "          9.315102e-01,  8.196849e-01],\n",
       "        [ 6.394557e-01,  6.162713e-01,  6.582525e-01,  6.849660e-01, ...,  7.615169e-01,  8.150862e-01,\n",
       "          7.435393e-01,  5.550404e-01]],\n",
       "\n",
       "       [[ 4.288197e-02,  3.586289e-02,  1.979762e-02,  4.157487e-02, ...,  1.990840e-01,  2.420095e-01,\n",
       "          2.877120e-01,  3.103569e-01],\n",
       "        [ 4.962894e-02,  3.502490e-02,  3.614373e-03,  7.033223e-05, ...,  1.563922e-01,  1.922820e-01,\n",
       "          2.452845e-01,  2.568748e-01],\n",
       "        [ 7.724187e-02,  4.813374e-02,  1.849236e-02, -1.296011e-02, ...,  1.404089e-01,  1.845620e-01,\n",
       "          2.243527e-01,  2.408762e-01],\n",
       "        [ 1.300122e-01,  1.025737e-01,  8.795914e-02,  4.576286e-02, ...,  1.721860e-01,  2.096085e-01,\n",
       "          2.493705e-01,  2.706239e-01],\n",
       "        ...,\n",
       "        [ 2.466994e-01,  1.628171e-01,  1.162946e-01,  1.563058e-01, ..., -6.753094e-02, -9.556959e-02,\n",
       "         -4.442038e-02, -9.396417e-03],\n",
       "        [ 1.840173e-01,  1.075276e-01,  6.186185e-02,  1.495719e-01, ...,  2.777258e-02, -3.337731e-02,\n",
       "         -5.033160e-02, -7.935677e-02],\n",
       "        [ 6.584295e-02,  8.974665e-02,  8.604161e-02,  9.996101e-02, ...,  2.057362e-01,  2.437397e-01,\n",
       "          2.530226e-01,  2.004631e-01],\n",
       "        [ 9.364849e-02,  2.017706e-01,  2.386460e-01,  1.540475e-01, ...,  2.210179e-01,  2.937904e-01,\n",
       "          3.533009e-01,  3.337062e-01]],\n",
       "\n",
       "       [[ 1.175666e+00,  1.220566e+00,  1.200156e+00,  1.138659e+00, ...,  2.616597e-01,  1.827675e-02,\n",
       "         -1.212664e-01, -2.315756e-01],\n",
       "        [ 1.205749e+00,  1.243035e+00,  1.218173e+00,  1.190289e+00, ...,  2.978287e-01,  3.892778e-02,\n",
       "         -6.948082e-02, -1.628896e-01],\n",
       "        [ 1.199235e+00,  1.229218e+00,  1.232467e+00,  1.257338e+00, ...,  3.268813e-01,  1.160583e-01,\n",
       "          4.168683e-03, -6.838661e-02],\n",
       "        [ 1.137234e+00,  1.183527e+00,  1.212620e+00,  1.240344e+00, ...,  3.642857e-01,  1.664781e-01,\n",
       "          6.743427e-02,  5.027659e-03],\n",
       "        ...,\n",
       "        [-6.436908e-01, -5.503952e-01, -4.742598e-01, -4.875617e-01, ..., -1.037436e+00, -1.040613e+00,\n",
       "         -1.177142e+00, -1.220379e+00],\n",
       "        [-7.436591e-01, -6.238915e-01, -4.735585e-01, -5.080217e-01, ..., -8.872296e-01, -7.897854e-01,\n",
       "         -9.088947e-01, -9.806532e-01],\n",
       "        [-6.747395e-01, -6.015897e-01, -4.995870e-01, -5.259740e-01, ..., -9.087684e-01, -8.189642e-01,\n",
       "         -9.191239e-01, -9.423386e-01],\n",
       "        [-5.969926e-01, -6.313480e-01, -6.231819e-01, -5.786868e-01, ..., -9.600114e-01, -9.989499e-01,\n",
       "         -1.157557e+00, -1.163065e+00]],\n",
       "\n",
       "       [[-6.557240e-01, -6.066412e-01, -6.021547e-01, -5.794508e-01, ..., -4.527237e-01, -4.429210e-01,\n",
       "         -4.480022e-01, -4.591919e-01],\n",
       "        [-6.415200e-01, -6.098990e-01, -5.912560e-01, -5.577772e-01, ..., -4.720501e-01, -4.568951e-01,\n",
       "         -4.618531e-01, -4.677817e-01],\n",
       "        [-6.225531e-01, -6.006425e-01, -5.806381e-01, -5.501241e-01, ..., -4.918782e-01, -4.865148e-01,\n",
       "         -4.856077e-01, -4.998765e-01],\n",
       "        [-6.047767e-01, -5.906608e-01, -5.885634e-01, -5.494895e-01, ..., -5.143436e-01, -5.171611e-01,\n",
       "         -5.178832e-01, -5.292746e-01],\n",
       "        ...,\n",
       "        [-1.869220e-02,  2.220134e-02, -2.010003e-02, -5.067986e-02, ...,  7.476944e-02,  7.811080e-02,\n",
       "          5.143915e-02,  3.505992e-02],\n",
       "        [ 4.170997e-03,  2.213846e-02, -1.570337e-02, -3.041676e-02, ...,  7.305323e-02,  8.682193e-02,\n",
       "          4.202233e-02,  1.745484e-02],\n",
       "        [ 7.753043e-05, -2.853164e-02,  2.213888e-03,  1.826694e-02, ..., -1.318107e-03,  2.810498e-02,\n",
       "         -3.546954e-02, -7.692661e-02],\n",
       "        [-3.465827e-02, -6.448118e-02,  1.741178e-02,  5.688276e-02, ...,  2.353423e-02,  5.509750e-02,\n",
       "         -2.259056e-02, -8.520117e-02]]], dtype=float32)"
      ]
     },
     "execution_count": 166,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch_output[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### pytorch to onnx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 167,
   "metadata": {},
   "outputs": [],
   "source": [
    "onnx_output = torch.onnx._export(model, torch_input, MODELS/f'{arch_name}.onnx', verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 168,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "graph torch-jit-export (\n",
      "  %0[FLOAT, 1x224x224x3]\n",
      ") initializers (\n",
      "  %1[FLOAT, 64x3x3x3]\n",
      "  %2[FLOAT, 64]\n",
      ") {\n",
      "  %3 = Transpose[perm = [0, 3, 1, 2]](%0)\n",
      "  %4 = Conv[dilations = [1, 1], group = 1, kernel_shape = [3, 3], pads = [0, 0, 0, 0], strides = [1, 1]](%3, %1, %2)\n",
      "  return %4\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "# Load the ONNX model\n",
    "model = onnx.load(MODELS/f\"{arch_name}.onnx\")\n",
    "\n",
    "# Check that the IR is well formed\n",
    "onnx.checker.check_model(model)\n",
    "\n",
    "# Print a human readable representation of the graph\n",
    "print(onnx.helper.printable_graph(model.graph))\n",
    "\n",
    "onnx_output = to_np(onnx_output)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 169,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[ 5.928250e-01,  6.495560e-01,  6.586254e-01,  6.328704e-01, ...,  1.145797e-01, -1.446192e-02,\n",
       "         -9.115485e-02, -1.700499e-01],\n",
       "        [ 6.237740e-01,  6.704966e-01,  6.704947e-01,  6.700016e-01, ...,  1.083156e-01, -3.941596e-02,\n",
       "         -9.333248e-02, -1.599272e-01],\n",
       "        [ 6.543528e-01,  7.007562e-01,  6.936352e-01,  7.207102e-01, ...,  1.495145e-01, -5.462829e-04,\n",
       "         -5.724194e-02, -1.157975e-01],\n",
       "        [ 6.586937e-01,  6.734817e-01,  6.918388e-01,  7.479682e-01, ...,  1.761432e-01,  5.087249e-02,\n",
       "         -1.233188e-02, -3.808299e-02],\n",
       "        ...,\n",
       "        [-5.286742e-01, -3.994761e-01, -4.185633e-01, -5.430715e-01, ..., -8.238600e-01, -8.971226e-01,\n",
       "         -1.071087e+00, -1.090646e+00],\n",
       "        [-6.421031e-01, -5.654278e-01, -4.625924e-01, -5.476409e-01, ..., -8.586741e-01, -7.533782e-01,\n",
       "         -8.619232e-01, -9.086895e-01],\n",
       "        [-6.917681e-01, -6.695323e-01, -5.033025e-01, -4.304609e-01, ..., -8.144889e-01, -6.926728e-01,\n",
       "         -7.351253e-01, -7.078845e-01],\n",
       "        [-5.975550e-01, -6.254424e-01, -5.287780e-01, -4.007401e-01, ..., -7.700869e-01, -7.775310e-01,\n",
       "         -9.107289e-01, -8.765057e-01]],\n",
       "\n",
       "       [[ 9.304108e-01,  9.061006e-01,  8.904200e-01,  8.748452e-01, ...,  1.507193e-01,  4.863602e-02,\n",
       "         -2.965286e-02, -6.887357e-02],\n",
       "        [ 9.354134e-01,  9.211414e-01,  9.025218e-01,  8.819546e-01, ...,  1.583953e-01,  7.783134e-02,\n",
       "          2.384748e-03, -2.762413e-02],\n",
       "        [ 9.247492e-01,  9.207308e-01,  9.197693e-01,  9.268373e-01, ...,  1.964801e-01,  1.070963e-01,\n",
       "          6.585616e-02,  3.885203e-02],\n",
       "        [ 8.749162e-01,  9.000026e-01,  9.111722e-01,  9.365388e-01, ...,  2.133995e-01,  1.541420e-01,\n",
       "          1.271267e-01,  1.043288e-01],\n",
       "        ...,\n",
       "        [-4.185115e-01, -4.692028e-01, -4.844882e-01, -4.011375e-01, ..., -9.281288e-01, -9.043365e-01,\n",
       "         -8.730649e-01, -8.296068e-01],\n",
       "        [-4.699255e-01, -4.824458e-01, -4.375248e-01, -3.610857e-01, ..., -6.990985e-01, -7.424051e-01,\n",
       "         -7.382967e-01, -6.652063e-01],\n",
       "        [-5.109918e-01, -4.493734e-01, -4.005437e-01, -4.306577e-01, ..., -6.778324e-01, -7.774227e-01,\n",
       "         -8.163334e-01, -7.388382e-01],\n",
       "        [-4.955843e-01, -4.066808e-01, -4.318877e-01, -5.693138e-01, ..., -6.130958e-01, -7.152777e-01,\n",
       "         -7.512010e-01, -6.329308e-01]],\n",
       "\n",
       "       [[-8.835173e-01, -8.811719e-01, -8.571362e-01, -8.381184e-01, ..., -1.226337e-01, -1.413189e-02,\n",
       "          6.252305e-02,  1.186951e-01],\n",
       "        [-8.997300e-01, -9.013266e-01, -8.836781e-01, -8.679811e-01, ..., -1.311588e-01, -3.768335e-02,\n",
       "          2.365048e-02,  6.502268e-02],\n",
       "        [-8.955600e-01, -9.047340e-01, -9.121482e-01, -9.156053e-01, ..., -1.767565e-01, -9.044454e-02,\n",
       "         -3.904480e-02, -1.396201e-03],\n",
       "        [-8.667954e-01, -8.885559e-01, -9.212384e-01, -9.478627e-01, ..., -2.214962e-01, -1.539418e-01,\n",
       "         -1.061586e-01, -7.190461e-02],\n",
       "        ...,\n",
       "        [ 6.527845e-01,  6.506559e-01,  6.483269e-01,  6.090642e-01, ...,  1.127803e+00,  1.157625e+00,\n",
       "          1.176368e+00,  1.130930e+00],\n",
       "        [ 7.203699e-01,  6.666610e-01,  6.309564e-01,  6.287088e-01, ...,  9.141809e-01,  9.502071e-01,\n",
       "          9.911880e-01,  9.432814e-01],\n",
       "        [ 7.402143e-01,  6.619299e-01,  6.260362e-01,  6.843429e-01, ...,  9.217322e-01,  9.615628e-01,\n",
       "          1.000422e+00,  9.200149e-01],\n",
       "        [ 7.193744e-01,  6.757457e-01,  6.706685e-01,  7.347934e-01, ...,  9.194486e-01,  9.976853e-01,\n",
       "          1.031135e+00,  9.027116e-01]],\n",
       "\n",
       "       [[-5.611107e-01, -6.185374e-01, -6.274537e-01, -6.018535e-01, ..., -3.286290e-01, -1.958319e-01,\n",
       "         -1.233654e-01, -5.534839e-02],\n",
       "        [-5.882207e-01, -6.269063e-01, -6.223419e-01, -5.886477e-01, ..., -2.825735e-01, -1.406765e-01,\n",
       "         -6.240558e-02, -1.910090e-03],\n",
       "        [-6.440452e-01, -6.522714e-01, -6.412833e-01, -5.956016e-01, ..., -2.842835e-01, -1.286600e-01,\n",
       "         -5.839922e-02, -1.226380e-02],\n",
       "        [-6.646396e-01, -6.643090e-01, -6.534098e-01, -6.352038e-01, ..., -3.086022e-01, -1.643063e-01,\n",
       "         -1.076260e-01, -7.506974e-02],\n",
       "        ...,\n",
       "        [ 1.974958e-01,  1.334879e-01,  2.476262e-01,  3.490084e-01, ...,  9.260128e-01,  9.112363e-01,\n",
       "          8.825827e-01,  8.980755e-01],\n",
       "        [ 2.174751e-01,  2.777752e-01,  3.892178e-01,  3.058348e-01, ...,  6.894872e-01,  6.306827e-01,\n",
       "          6.324791e-01,  7.319039e-01],\n",
       "        [ 4.317094e-01,  5.105460e-01,  4.332957e-01,  2.208417e-01, ...,  6.646112e-01,  5.065519e-01,\n",
       "          4.617480e-01,  6.012140e-01],\n",
       "        [ 4.610324e-01,  4.142273e-01,  2.067479e-01,  1.124400e-01, ...,  3.580654e-01,  2.135994e-01,\n",
       "          1.807216e-01,  3.309866e-01]],\n",
       "\n",
       "       ...,\n",
       "\n",
       "       [[-1.021994e+00, -1.007560e+00, -1.020993e+00, -1.028300e+00, ..., -6.495625e-01, -6.656052e-01,\n",
       "         -6.676931e-01, -6.754252e-01],\n",
       "        [-1.022150e+00, -1.020793e+00, -1.031548e+00, -1.039426e+00, ..., -6.874483e-01, -6.935481e-01,\n",
       "         -7.095402e-01, -7.143820e-01],\n",
       "        [-1.000767e+00, -1.022446e+00, -1.045676e+00, -1.072399e+00, ..., -7.137644e-01, -7.332122e-01,\n",
       "         -7.495062e-01, -7.635692e-01],\n",
       "        [-9.922659e-01, -1.020445e+00, -1.055791e+00, -1.083923e+00, ..., -7.817140e-01, -7.956998e-01,\n",
       "         -8.141065e-01, -8.214797e-01],\n",
       "        ...,\n",
       "        [ 5.825371e-01,  5.784728e-01,  5.523584e-01,  5.299688e-01, ...,  1.002575e+00,  1.063073e+00,\n",
       "          1.038671e+00,  9.423774e-01],\n",
       "        [ 5.757274e-01,  5.073735e-01,  4.978513e-01,  5.697755e-01, ...,  7.133898e-01,  8.178766e-01,\n",
       "          8.427096e-01,  7.569195e-01],\n",
       "        [ 6.112009e-01,  5.420885e-01,  5.843502e-01,  6.813412e-01, ...,  8.724557e-01,  9.414760e-01,\n",
       "          9.315102e-01,  8.196849e-01],\n",
       "        [ 6.394557e-01,  6.162713e-01,  6.582525e-01,  6.849660e-01, ...,  7.615169e-01,  8.150862e-01,\n",
       "          7.435393e-01,  5.550404e-01]],\n",
       "\n",
       "       [[ 4.288197e-02,  3.586289e-02,  1.979762e-02,  4.157487e-02, ...,  1.990840e-01,  2.420095e-01,\n",
       "          2.877120e-01,  3.103569e-01],\n",
       "        [ 4.962894e-02,  3.502490e-02,  3.614373e-03,  7.033223e-05, ...,  1.563922e-01,  1.922820e-01,\n",
       "          2.452845e-01,  2.568748e-01],\n",
       "        [ 7.724187e-02,  4.813374e-02,  1.849236e-02, -1.296011e-02, ...,  1.404089e-01,  1.845620e-01,\n",
       "          2.243527e-01,  2.408762e-01],\n",
       "        [ 1.300122e-01,  1.025737e-01,  8.795914e-02,  4.576286e-02, ...,  1.721860e-01,  2.096085e-01,\n",
       "          2.493705e-01,  2.706239e-01],\n",
       "        ...,\n",
       "        [ 2.466994e-01,  1.628171e-01,  1.162946e-01,  1.563058e-01, ..., -6.753094e-02, -9.556959e-02,\n",
       "         -4.442038e-02, -9.396417e-03],\n",
       "        [ 1.840173e-01,  1.075276e-01,  6.186185e-02,  1.495719e-01, ...,  2.777258e-02, -3.337731e-02,\n",
       "         -5.033160e-02, -7.935677e-02],\n",
       "        [ 6.584295e-02,  8.974665e-02,  8.604161e-02,  9.996101e-02, ...,  2.057362e-01,  2.437397e-01,\n",
       "          2.530226e-01,  2.004631e-01],\n",
       "        [ 9.364849e-02,  2.017706e-01,  2.386460e-01,  1.540475e-01, ...,  2.210179e-01,  2.937904e-01,\n",
       "          3.533009e-01,  3.337062e-01]],\n",
       "\n",
       "       [[ 1.175666e+00,  1.220566e+00,  1.200156e+00,  1.138659e+00, ...,  2.616597e-01,  1.827675e-02,\n",
       "         -1.212664e-01, -2.315756e-01],\n",
       "        [ 1.205749e+00,  1.243035e+00,  1.218173e+00,  1.190289e+00, ...,  2.978287e-01,  3.892778e-02,\n",
       "         -6.948082e-02, -1.628896e-01],\n",
       "        [ 1.199235e+00,  1.229218e+00,  1.232467e+00,  1.257338e+00, ...,  3.268813e-01,  1.160583e-01,\n",
       "          4.168683e-03, -6.838661e-02],\n",
       "        [ 1.137234e+00,  1.183527e+00,  1.212620e+00,  1.240344e+00, ...,  3.642857e-01,  1.664781e-01,\n",
       "          6.743427e-02,  5.027659e-03],\n",
       "        ...,\n",
       "        [-6.436908e-01, -5.503952e-01, -4.742598e-01, -4.875617e-01, ..., -1.037436e+00, -1.040613e+00,\n",
       "         -1.177142e+00, -1.220379e+00],\n",
       "        [-7.436591e-01, -6.238915e-01, -4.735585e-01, -5.080217e-01, ..., -8.872296e-01, -7.897854e-01,\n",
       "         -9.088947e-01, -9.806532e-01],\n",
       "        [-6.747395e-01, -6.015897e-01, -4.995870e-01, -5.259740e-01, ..., -9.087684e-01, -8.189642e-01,\n",
       "         -9.191239e-01, -9.423386e-01],\n",
       "        [-5.969926e-01, -6.313480e-01, -6.231819e-01, -5.786868e-01, ..., -9.600114e-01, -9.989499e-01,\n",
       "         -1.157557e+00, -1.163065e+00]],\n",
       "\n",
       "       [[-6.557240e-01, -6.066412e-01, -6.021547e-01, -5.794508e-01, ..., -4.527237e-01, -4.429210e-01,\n",
       "         -4.480022e-01, -4.591919e-01],\n",
       "        [-6.415200e-01, -6.098990e-01, -5.912560e-01, -5.577772e-01, ..., -4.720501e-01, -4.568951e-01,\n",
       "         -4.618531e-01, -4.677817e-01],\n",
       "        [-6.225531e-01, -6.006425e-01, -5.806381e-01, -5.501241e-01, ..., -4.918782e-01, -4.865148e-01,\n",
       "         -4.856077e-01, -4.998765e-01],\n",
       "        [-6.047767e-01, -5.906608e-01, -5.885634e-01, -5.494895e-01, ..., -5.143436e-01, -5.171611e-01,\n",
       "         -5.178832e-01, -5.292746e-01],\n",
       "        ...,\n",
       "        [-1.869220e-02,  2.220134e-02, -2.010003e-02, -5.067986e-02, ...,  7.476944e-02,  7.811080e-02,\n",
       "          5.143915e-02,  3.505992e-02],\n",
       "        [ 4.170997e-03,  2.213846e-02, -1.570337e-02, -3.041676e-02, ...,  7.305323e-02,  8.682193e-02,\n",
       "          4.202233e-02,  1.745484e-02],\n",
       "        [ 7.753043e-05, -2.853164e-02,  2.213888e-03,  1.826694e-02, ..., -1.318107e-03,  2.810498e-02,\n",
       "         -3.546954e-02, -7.692661e-02],\n",
       "        [-3.465827e-02, -6.448118e-02,  1.741178e-02,  5.688276e-02, ...,  2.353423e-02,  5.509750e-02,\n",
       "         -2.259056e-02, -8.520117e-02]]], dtype=float32)"
      ]
     },
     "execution_count": 169,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch_output[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 170,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[ 5.928250e-01,  6.495560e-01,  6.586254e-01,  6.328704e-01, ...,  1.145797e-01, -1.446192e-02,\n",
       "         -9.115485e-02, -1.700499e-01],\n",
       "        [ 6.237740e-01,  6.704966e-01,  6.704947e-01,  6.700016e-01, ...,  1.083156e-01, -3.941596e-02,\n",
       "         -9.333248e-02, -1.599272e-01],\n",
       "        [ 6.543528e-01,  7.007562e-01,  6.936352e-01,  7.207102e-01, ...,  1.495145e-01, -5.462829e-04,\n",
       "         -5.724194e-02, -1.157975e-01],\n",
       "        [ 6.586937e-01,  6.734817e-01,  6.918388e-01,  7.479682e-01, ...,  1.761432e-01,  5.087249e-02,\n",
       "         -1.233188e-02, -3.808299e-02],\n",
       "        ...,\n",
       "        [-5.286742e-01, -3.994761e-01, -4.185633e-01, -5.430715e-01, ..., -8.238600e-01, -8.971226e-01,\n",
       "         -1.071087e+00, -1.090646e+00],\n",
       "        [-6.421031e-01, -5.654278e-01, -4.625924e-01, -5.476409e-01, ..., -8.586741e-01, -7.533782e-01,\n",
       "         -8.619232e-01, -9.086895e-01],\n",
       "        [-6.917681e-01, -6.695323e-01, -5.033025e-01, -4.304609e-01, ..., -8.144889e-01, -6.926728e-01,\n",
       "         -7.351253e-01, -7.078845e-01],\n",
       "        [-5.975550e-01, -6.254424e-01, -5.287780e-01, -4.007401e-01, ..., -7.700869e-01, -7.775310e-01,\n",
       "         -9.107289e-01, -8.765057e-01]],\n",
       "\n",
       "       [[ 9.304108e-01,  9.061006e-01,  8.904200e-01,  8.748452e-01, ...,  1.507193e-01,  4.863602e-02,\n",
       "         -2.965286e-02, -6.887357e-02],\n",
       "        [ 9.354134e-01,  9.211414e-01,  9.025218e-01,  8.819546e-01, ...,  1.583953e-01,  7.783134e-02,\n",
       "          2.384748e-03, -2.762413e-02],\n",
       "        [ 9.247492e-01,  9.207308e-01,  9.197693e-01,  9.268373e-01, ...,  1.964801e-01,  1.070963e-01,\n",
       "          6.585616e-02,  3.885203e-02],\n",
       "        [ 8.749162e-01,  9.000026e-01,  9.111722e-01,  9.365388e-01, ...,  2.133995e-01,  1.541420e-01,\n",
       "          1.271267e-01,  1.043288e-01],\n",
       "        ...,\n",
       "        [-4.185115e-01, -4.692028e-01, -4.844882e-01, -4.011375e-01, ..., -9.281288e-01, -9.043365e-01,\n",
       "         -8.730649e-01, -8.296068e-01],\n",
       "        [-4.699255e-01, -4.824458e-01, -4.375248e-01, -3.610857e-01, ..., -6.990985e-01, -7.424051e-01,\n",
       "         -7.382967e-01, -6.652063e-01],\n",
       "        [-5.109918e-01, -4.493734e-01, -4.005437e-01, -4.306577e-01, ..., -6.778324e-01, -7.774227e-01,\n",
       "         -8.163334e-01, -7.388382e-01],\n",
       "        [-4.955843e-01, -4.066808e-01, -4.318877e-01, -5.693138e-01, ..., -6.130958e-01, -7.152777e-01,\n",
       "         -7.512010e-01, -6.329308e-01]],\n",
       "\n",
       "       [[-8.835173e-01, -8.811719e-01, -8.571362e-01, -8.381184e-01, ..., -1.226337e-01, -1.413189e-02,\n",
       "          6.252305e-02,  1.186951e-01],\n",
       "        [-8.997300e-01, -9.013266e-01, -8.836781e-01, -8.679811e-01, ..., -1.311588e-01, -3.768335e-02,\n",
       "          2.365048e-02,  6.502268e-02],\n",
       "        [-8.955600e-01, -9.047340e-01, -9.121482e-01, -9.156053e-01, ..., -1.767565e-01, -9.044454e-02,\n",
       "         -3.904480e-02, -1.396201e-03],\n",
       "        [-8.667954e-01, -8.885559e-01, -9.212384e-01, -9.478627e-01, ..., -2.214962e-01, -1.539418e-01,\n",
       "         -1.061586e-01, -7.190461e-02],\n",
       "        ...,\n",
       "        [ 6.527845e-01,  6.506559e-01,  6.483269e-01,  6.090642e-01, ...,  1.127803e+00,  1.157625e+00,\n",
       "          1.176368e+00,  1.130930e+00],\n",
       "        [ 7.203699e-01,  6.666610e-01,  6.309564e-01,  6.287088e-01, ...,  9.141809e-01,  9.502071e-01,\n",
       "          9.911880e-01,  9.432814e-01],\n",
       "        [ 7.402143e-01,  6.619299e-01,  6.260362e-01,  6.843429e-01, ...,  9.217322e-01,  9.615628e-01,\n",
       "          1.000422e+00,  9.200149e-01],\n",
       "        [ 7.193744e-01,  6.757457e-01,  6.706685e-01,  7.347934e-01, ...,  9.194486e-01,  9.976853e-01,\n",
       "          1.031135e+00,  9.027116e-01]],\n",
       "\n",
       "       [[-5.611107e-01, -6.185374e-01, -6.274537e-01, -6.018535e-01, ..., -3.286290e-01, -1.958319e-01,\n",
       "         -1.233654e-01, -5.534839e-02],\n",
       "        [-5.882207e-01, -6.269063e-01, -6.223419e-01, -5.886477e-01, ..., -2.825735e-01, -1.406765e-01,\n",
       "         -6.240558e-02, -1.910090e-03],\n",
       "        [-6.440452e-01, -6.522714e-01, -6.412833e-01, -5.956016e-01, ..., -2.842835e-01, -1.286600e-01,\n",
       "         -5.839922e-02, -1.226380e-02],\n",
       "        [-6.646396e-01, -6.643090e-01, -6.534098e-01, -6.352038e-01, ..., -3.086022e-01, -1.643063e-01,\n",
       "         -1.076260e-01, -7.506974e-02],\n",
       "        ...,\n",
       "        [ 1.974958e-01,  1.334879e-01,  2.476262e-01,  3.490084e-01, ...,  9.260128e-01,  9.112363e-01,\n",
       "          8.825827e-01,  8.980755e-01],\n",
       "        [ 2.174751e-01,  2.777752e-01,  3.892178e-01,  3.058348e-01, ...,  6.894872e-01,  6.306827e-01,\n",
       "          6.324791e-01,  7.319039e-01],\n",
       "        [ 4.317094e-01,  5.105460e-01,  4.332957e-01,  2.208417e-01, ...,  6.646112e-01,  5.065519e-01,\n",
       "          4.617480e-01,  6.012140e-01],\n",
       "        [ 4.610324e-01,  4.142273e-01,  2.067479e-01,  1.124400e-01, ...,  3.580654e-01,  2.135994e-01,\n",
       "          1.807216e-01,  3.309866e-01]],\n",
       "\n",
       "       ...,\n",
       "\n",
       "       [[-1.021994e+00, -1.007560e+00, -1.020993e+00, -1.028300e+00, ..., -6.495625e-01, -6.656052e-01,\n",
       "         -6.676931e-01, -6.754252e-01],\n",
       "        [-1.022150e+00, -1.020793e+00, -1.031548e+00, -1.039426e+00, ..., -6.874483e-01, -6.935481e-01,\n",
       "         -7.095402e-01, -7.143820e-01],\n",
       "        [-1.000767e+00, -1.022446e+00, -1.045676e+00, -1.072399e+00, ..., -7.137644e-01, -7.332122e-01,\n",
       "         -7.495062e-01, -7.635692e-01],\n",
       "        [-9.922659e-01, -1.020445e+00, -1.055791e+00, -1.083923e+00, ..., -7.817140e-01, -7.956998e-01,\n",
       "         -8.141065e-01, -8.214797e-01],\n",
       "        ...,\n",
       "        [ 5.825371e-01,  5.784728e-01,  5.523584e-01,  5.299688e-01, ...,  1.002575e+00,  1.063073e+00,\n",
       "          1.038671e+00,  9.423774e-01],\n",
       "        [ 5.757274e-01,  5.073735e-01,  4.978513e-01,  5.697755e-01, ...,  7.133898e-01,  8.178766e-01,\n",
       "          8.427096e-01,  7.569195e-01],\n",
       "        [ 6.112009e-01,  5.420885e-01,  5.843502e-01,  6.813412e-01, ...,  8.724557e-01,  9.414760e-01,\n",
       "          9.315102e-01,  8.196849e-01],\n",
       "        [ 6.394557e-01,  6.162713e-01,  6.582525e-01,  6.849660e-01, ...,  7.615169e-01,  8.150862e-01,\n",
       "          7.435393e-01,  5.550404e-01]],\n",
       "\n",
       "       [[ 4.288197e-02,  3.586289e-02,  1.979762e-02,  4.157487e-02, ...,  1.990840e-01,  2.420095e-01,\n",
       "          2.877120e-01,  3.103569e-01],\n",
       "        [ 4.962894e-02,  3.502490e-02,  3.614373e-03,  7.033223e-05, ...,  1.563922e-01,  1.922820e-01,\n",
       "          2.452845e-01,  2.568748e-01],\n",
       "        [ 7.724187e-02,  4.813374e-02,  1.849236e-02, -1.296011e-02, ...,  1.404089e-01,  1.845620e-01,\n",
       "          2.243527e-01,  2.408762e-01],\n",
       "        [ 1.300122e-01,  1.025737e-01,  8.795914e-02,  4.576286e-02, ...,  1.721860e-01,  2.096085e-01,\n",
       "          2.493705e-01,  2.706239e-01],\n",
       "        ...,\n",
       "        [ 2.466994e-01,  1.628171e-01,  1.162946e-01,  1.563058e-01, ..., -6.753094e-02, -9.556959e-02,\n",
       "         -4.442038e-02, -9.396417e-03],\n",
       "        [ 1.840173e-01,  1.075276e-01,  6.186185e-02,  1.495719e-01, ...,  2.777258e-02, -3.337731e-02,\n",
       "         -5.033160e-02, -7.935677e-02],\n",
       "        [ 6.584295e-02,  8.974665e-02,  8.604161e-02,  9.996101e-02, ...,  2.057362e-01,  2.437397e-01,\n",
       "          2.530226e-01,  2.004631e-01],\n",
       "        [ 9.364849e-02,  2.017706e-01,  2.386460e-01,  1.540475e-01, ...,  2.210179e-01,  2.937904e-01,\n",
       "          3.533009e-01,  3.337062e-01]],\n",
       "\n",
       "       [[ 1.175666e+00,  1.220566e+00,  1.200156e+00,  1.138659e+00, ...,  2.616597e-01,  1.827675e-02,\n",
       "         -1.212664e-01, -2.315756e-01],\n",
       "        [ 1.205749e+00,  1.243035e+00,  1.218173e+00,  1.190289e+00, ...,  2.978287e-01,  3.892778e-02,\n",
       "         -6.948082e-02, -1.628896e-01],\n",
       "        [ 1.199235e+00,  1.229218e+00,  1.232467e+00,  1.257338e+00, ...,  3.268813e-01,  1.160583e-01,\n",
       "          4.168683e-03, -6.838661e-02],\n",
       "        [ 1.137234e+00,  1.183527e+00,  1.212620e+00,  1.240344e+00, ...,  3.642857e-01,  1.664781e-01,\n",
       "          6.743427e-02,  5.027659e-03],\n",
       "        ...,\n",
       "        [-6.436908e-01, -5.503952e-01, -4.742598e-01, -4.875617e-01, ..., -1.037436e+00, -1.040613e+00,\n",
       "         -1.177142e+00, -1.220379e+00],\n",
       "        [-7.436591e-01, -6.238915e-01, -4.735585e-01, -5.080217e-01, ..., -8.872296e-01, -7.897854e-01,\n",
       "         -9.088947e-01, -9.806532e-01],\n",
       "        [-6.747395e-01, -6.015897e-01, -4.995870e-01, -5.259740e-01, ..., -9.087684e-01, -8.189642e-01,\n",
       "         -9.191239e-01, -9.423386e-01],\n",
       "        [-5.969926e-01, -6.313480e-01, -6.231819e-01, -5.786868e-01, ..., -9.600114e-01, -9.989499e-01,\n",
       "         -1.157557e+00, -1.163065e+00]],\n",
       "\n",
       "       [[-6.557240e-01, -6.066412e-01, -6.021547e-01, -5.794508e-01, ..., -4.527237e-01, -4.429210e-01,\n",
       "         -4.480022e-01, -4.591919e-01],\n",
       "        [-6.415200e-01, -6.098990e-01, -5.912560e-01, -5.577772e-01, ..., -4.720501e-01, -4.568951e-01,\n",
       "         -4.618531e-01, -4.677817e-01],\n",
       "        [-6.225531e-01, -6.006425e-01, -5.806381e-01, -5.501241e-01, ..., -4.918782e-01, -4.865148e-01,\n",
       "         -4.856077e-01, -4.998765e-01],\n",
       "        [-6.047767e-01, -5.906608e-01, -5.885634e-01, -5.494895e-01, ..., -5.143436e-01, -5.171611e-01,\n",
       "         -5.178832e-01, -5.292746e-01],\n",
       "        ...,\n",
       "        [-1.869220e-02,  2.220134e-02, -2.010003e-02, -5.067986e-02, ...,  7.476944e-02,  7.811080e-02,\n",
       "          5.143915e-02,  3.505992e-02],\n",
       "        [ 4.170997e-03,  2.213846e-02, -1.570337e-02, -3.041676e-02, ...,  7.305323e-02,  8.682193e-02,\n",
       "          4.202233e-02,  1.745484e-02],\n",
       "        [ 7.753043e-05, -2.853164e-02,  2.213888e-03,  1.826694e-02, ..., -1.318107e-03,  2.810498e-02,\n",
       "         -3.546954e-02, -7.692661e-02],\n",
       "        [-3.465827e-02, -6.448118e-02,  1.741178e-02,  5.688276e-02, ...,  2.353423e-02,  5.509750e-02,\n",
       "         -2.259056e-02, -8.520117e-02]]], dtype=float32)"
      ]
     },
     "execution_count": 170,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "onnx_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 171,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not inputs_same(torch_output[0], onnx_output):\n",
    "    raise Exception(\"PyTorch and ONNX outputs are not same\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### onnx to coreml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 172,
   "metadata": {},
   "outputs": [],
   "source": [
    "from onnx_coreml import convert"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 193,
   "metadata": {},
   "outputs": [],
   "source": [
    "import onnx_coreml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 173,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1/2: Converting Node Type Transpose\n",
      "2/2: Converting Node Type Conv\n",
      "Translation to CoreML spec completed. Now compiling the CoreML model.\n",
      "Model Compilation done.\n"
     ]
    }
   ],
   "source": [
    "# convert to coreml\n",
    "model = onnx.load(MODELS/f\"{arch_name}.onnx\")\n",
    "mlmodel = convert(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 174,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlmodel.save(MODELS/f'{arch_name}.mlmodel')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 175,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Exception: Model prediction is only supported on macOS version 10.13 or later.\n",
    "# torch with coreml - only on mac OS\n",
    "# mlmodel.predict({\"0\":numpy_input})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### onnx to tf \n",
    "\n",
    "You should get an warning message says\n",
    "\n",
    "UserWarning: Using the pooling op in compatibility mode.This means your graph cannot be serialized.\n",
    "Please configure your pooling operation to only use paddings that correspond to Tensorflow SAME or VALID padding.\n",
    "One layer in your network can not use native tensorflow op. We have to use compatible pool.\n",
    "I checked and it seems following layer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 176,
   "metadata": {},
   "outputs": [],
   "source": [
    "from onnx_tf.backend import prepare"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 177,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/onnx-tensorflow/onnx_tf/common/handler_helper.py:73: UserWarning: Unknown op ConstantFill in domain `ai.onnx`.\n",
      "  handler.ONNX_OP, handler.DOMAIN or \"ai.onnx\"))\n",
      "/home/ubuntu/onnx-tensorflow/onnx_tf/common/handler_helper.py:73: UserWarning: Unknown op MeanVarianceNormalization in domain `ai.onnx`.\n",
      "  handler.ONNX_OP, handler.DOMAIN or \"ai.onnx\"))\n"
     ]
    }
   ],
   "source": [
    "model = onnx.load(MODELS/f\"{arch_name}.onnx\")\n",
    "tf_rep = prepare(model)\n",
    "tf_output = tf_rep.run(numpy_input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 178,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<tf.Operation 'Const' type=Const>,\n",
       " <tf.Operation 'Const_1' type=Const>,\n",
       " <tf.Operation '0' type=Placeholder>,\n",
       " <tf.Operation 'transpose/perm' type=Const>,\n",
       " <tf.Operation 'transpose' type=Transpose>,\n",
       " <tf.Operation 'transpose_1/perm' type=Const>,\n",
       " <tf.Operation 'transpose_1' type=Transpose>,\n",
       " <tf.Operation 'Const_2' type=Const>,\n",
       " <tf.Operation 'Pad' type=Pad>,\n",
       " <tf.Operation 'Const_3' type=Const>,\n",
       " <tf.Operation 'split/split_dim' type=Const>,\n",
       " <tf.Operation 'split' type=Split>,\n",
       " <tf.Operation 'transpose_2/perm' type=Const>,\n",
       " <tf.Operation 'transpose_2' type=Transpose>,\n",
       " <tf.Operation 'Const_4' type=Const>,\n",
       " <tf.Operation 'split_1/split_dim' type=Const>,\n",
       " <tf.Operation 'split_1' type=Split>,\n",
       " <tf.Operation 'convolution/dilation_rate' type=Const>,\n",
       " <tf.Operation 'convolution' type=Conv2D>,\n",
       " <tf.Operation 'concat/concat_dim' type=Const>,\n",
       " <tf.Operation 'concat' type=Identity>,\n",
       " <tf.Operation 'Add' type=Add>,\n",
       " <tf.Operation 'transpose_3/perm' type=Const>,\n",
       " <tf.Operation 'transpose_3' type=Transpose>,\n",
       " <tf.Operation 'init' type=NoOp>]"
      ]
     },
     "execution_count": 178,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf_rep.graph.get_operations()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 179,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[ 5.928249e-01,  6.495559e-01,  6.586254e-01,  6.328702e-01, ...,  1.145797e-01, -1.446191e-02,\n",
       "         -9.115484e-02, -1.700499e-01],\n",
       "        [ 6.237740e-01,  6.704966e-01,  6.704946e-01,  6.700016e-01, ...,  1.083155e-01, -3.941594e-02,\n",
       "         -9.333252e-02, -1.599272e-01],\n",
       "        [ 6.543528e-01,  7.007561e-01,  6.936354e-01,  7.207103e-01, ...,  1.495145e-01, -5.463138e-04,\n",
       "         -5.724193e-02, -1.157975e-01],\n",
       "        [ 6.586937e-01,  6.734816e-01,  6.918389e-01,  7.479682e-01, ...,  1.761432e-01,  5.087247e-02,\n",
       "         -1.233190e-02, -3.808297e-02],\n",
       "        ...,\n",
       "        [-5.286742e-01, -3.994760e-01, -4.185633e-01, -5.430715e-01, ..., -8.238600e-01, -8.971226e-01,\n",
       "         -1.071087e+00, -1.090645e+00],\n",
       "        [-6.421031e-01, -5.654277e-01, -4.625923e-01, -5.476409e-01, ..., -8.586740e-01, -7.533783e-01,\n",
       "         -8.619232e-01, -9.086896e-01],\n",
       "        [-6.917681e-01, -6.695324e-01, -5.033025e-01, -4.304608e-01, ..., -8.144889e-01, -6.926728e-01,\n",
       "         -7.351254e-01, -7.078845e-01],\n",
       "        [-5.975550e-01, -6.254423e-01, -5.287780e-01, -4.007401e-01, ..., -7.700868e-01, -7.775310e-01,\n",
       "         -9.107289e-01, -8.765057e-01]],\n",
       "\n",
       "       [[ 9.304109e-01,  9.061006e-01,  8.904200e-01,  8.748451e-01, ...,  1.507193e-01,  4.863601e-02,\n",
       "         -2.965288e-02, -6.887355e-02],\n",
       "        [ 9.354135e-01,  9.211414e-01,  9.025218e-01,  8.819547e-01, ...,  1.583953e-01,  7.783135e-02,\n",
       "          2.384741e-03, -2.762413e-02],\n",
       "        [ 9.247494e-01,  9.207308e-01,  9.197693e-01,  9.268373e-01, ...,  1.964801e-01,  1.070963e-01,\n",
       "          6.585614e-02,  3.885203e-02],\n",
       "        [ 8.749162e-01,  9.000027e-01,  9.111722e-01,  9.365388e-01, ...,  2.133996e-01,  1.541421e-01,\n",
       "          1.271266e-01,  1.043288e-01],\n",
       "        ...,\n",
       "        [-4.185115e-01, -4.692028e-01, -4.844882e-01, -4.011376e-01, ..., -9.281287e-01, -9.043365e-01,\n",
       "         -8.730650e-01, -8.296068e-01],\n",
       "        [-4.699254e-01, -4.824459e-01, -4.375247e-01, -3.610858e-01, ..., -6.990983e-01, -7.424051e-01,\n",
       "         -7.382967e-01, -6.652063e-01],\n",
       "        [-5.109919e-01, -4.493735e-01, -4.005437e-01, -4.306577e-01, ..., -6.778324e-01, -7.774227e-01,\n",
       "         -8.163334e-01, -7.388381e-01],\n",
       "        [-4.955844e-01, -4.066808e-01, -4.318877e-01, -5.693137e-01, ..., -6.130958e-01, -7.152777e-01,\n",
       "         -7.512011e-01, -6.329307e-01]],\n",
       "\n",
       "       [[-8.835173e-01, -8.811720e-01, -8.571363e-01, -8.381183e-01, ..., -1.226337e-01, -1.413189e-02,\n",
       "          6.252309e-02,  1.186951e-01],\n",
       "        [-8.997299e-01, -9.013266e-01, -8.836781e-01, -8.679811e-01, ..., -1.311588e-01, -3.768332e-02,\n",
       "          2.365047e-02,  6.502274e-02],\n",
       "        [-8.955598e-01, -9.047340e-01, -9.121482e-01, -9.156054e-01, ..., -1.767565e-01, -9.044456e-02,\n",
       "         -3.904480e-02, -1.396231e-03],\n",
       "        [-8.667955e-01, -8.885559e-01, -9.212384e-01, -9.478628e-01, ..., -2.214962e-01, -1.539418e-01,\n",
       "         -1.061586e-01, -7.190464e-02],\n",
       "        ...,\n",
       "        [ 6.527844e-01,  6.506559e-01,  6.483268e-01,  6.090642e-01, ...,  1.127803e+00,  1.157625e+00,\n",
       "          1.176368e+00,  1.130930e+00],\n",
       "        [ 7.203699e-01,  6.666610e-01,  6.309563e-01,  6.287088e-01, ...,  9.141810e-01,  9.502069e-01,\n",
       "          9.911880e-01,  9.432812e-01],\n",
       "        [ 7.402142e-01,  6.619299e-01,  6.260363e-01,  6.843429e-01, ...,  9.217322e-01,  9.615628e-01,\n",
       "          1.000422e+00,  9.200149e-01],\n",
       "        [ 7.193743e-01,  6.757457e-01,  6.706684e-01,  7.347934e-01, ...,  9.194484e-01,  9.976852e-01,\n",
       "          1.031135e+00,  9.027117e-01]],\n",
       "\n",
       "       [[-5.611106e-01, -6.185372e-01, -6.274537e-01, -6.018535e-01, ..., -3.286291e-01, -1.958319e-01,\n",
       "         -1.233654e-01, -5.534839e-02],\n",
       "        [-5.882207e-01, -6.269063e-01, -6.223416e-01, -5.886477e-01, ..., -2.825735e-01, -1.406765e-01,\n",
       "         -6.240559e-02, -1.910051e-03],\n",
       "        [-6.440450e-01, -6.522713e-01, -6.412832e-01, -5.956013e-01, ..., -2.842835e-01, -1.286600e-01,\n",
       "         -5.839920e-02, -1.226379e-02],\n",
       "        [-6.646397e-01, -6.643088e-01, -6.534097e-01, -6.352037e-01, ..., -3.086021e-01, -1.643063e-01,\n",
       "         -1.076260e-01, -7.506977e-02],\n",
       "        ...,\n",
       "        [ 1.974958e-01,  1.334879e-01,  2.476263e-01,  3.490084e-01, ...,  9.260127e-01,  9.112365e-01,\n",
       "          8.825831e-01,  8.980753e-01],\n",
       "        [ 2.174751e-01,  2.777752e-01,  3.892177e-01,  3.058347e-01, ...,  6.894875e-01,  6.306829e-01,\n",
       "          6.324791e-01,  7.319040e-01],\n",
       "        [ 4.317095e-01,  5.105460e-01,  4.332959e-01,  2.208420e-01, ...,  6.646112e-01,  5.065519e-01,\n",
       "          4.617481e-01,  6.012139e-01],\n",
       "        [ 4.610323e-01,  4.142273e-01,  2.067479e-01,  1.124399e-01, ...,  3.580655e-01,  2.135992e-01,\n",
       "          1.807216e-01,  3.309866e-01]],\n",
       "\n",
       "       ...,\n",
       "\n",
       "       [[-1.021994e+00, -1.007560e+00, -1.020993e+00, -1.028299e+00, ..., -6.495626e-01, -6.656051e-01,\n",
       "         -6.676931e-01, -6.754251e-01],\n",
       "        [-1.022150e+00, -1.020793e+00, -1.031548e+00, -1.039426e+00, ..., -6.874483e-01, -6.935482e-01,\n",
       "         -7.095402e-01, -7.143820e-01],\n",
       "        [-1.000767e+00, -1.022446e+00, -1.045676e+00, -1.072399e+00, ..., -7.137645e-01, -7.332121e-01,\n",
       "         -7.495061e-01, -7.635693e-01],\n",
       "        [-9.922659e-01, -1.020445e+00, -1.055791e+00, -1.083923e+00, ..., -7.817139e-01, -7.956998e-01,\n",
       "         -8.141064e-01, -8.214797e-01],\n",
       "        ...,\n",
       "        [ 5.825371e-01,  5.784727e-01,  5.523583e-01,  5.299689e-01, ...,  1.002575e+00,  1.063073e+00,\n",
       "          1.038671e+00,  9.423774e-01],\n",
       "        [ 5.757274e-01,  5.073734e-01,  4.978513e-01,  5.697754e-01, ...,  7.133896e-01,  8.178766e-01,\n",
       "          8.427094e-01,  7.569194e-01],\n",
       "        [ 6.112007e-01,  5.420885e-01,  5.843502e-01,  6.813412e-01, ...,  8.724557e-01,  9.414759e-01,\n",
       "          9.315102e-01,  8.196850e-01],\n",
       "        [ 6.394557e-01,  6.162713e-01,  6.582525e-01,  6.849660e-01, ...,  7.615168e-01,  8.150862e-01,\n",
       "          7.435393e-01,  5.550404e-01]],\n",
       "\n",
       "       [[ 4.288198e-02,  3.586289e-02,  1.979771e-02,  4.157487e-02, ...,  1.990840e-01,  2.420095e-01,\n",
       "          2.877121e-01,  3.103569e-01],\n",
       "        [ 4.962896e-02,  3.502500e-02,  3.614321e-03,  7.040799e-05, ...,  1.563922e-01,  1.922820e-01,\n",
       "          2.452845e-01,  2.568747e-01],\n",
       "        [ 7.724193e-02,  4.813366e-02,  1.849239e-02, -1.296002e-02, ...,  1.404090e-01,  1.845620e-01,\n",
       "          2.243527e-01,  2.408761e-01],\n",
       "        [ 1.300122e-01,  1.025737e-01,  8.795923e-02,  4.576284e-02, ...,  1.721859e-01,  2.096085e-01,\n",
       "          2.493705e-01,  2.706239e-01],\n",
       "        ...,\n",
       "        [ 2.466994e-01,  1.628171e-01,  1.162947e-01,  1.563058e-01, ..., -6.753095e-02, -9.556955e-02,\n",
       "         -4.442039e-02, -9.396434e-03],\n",
       "        [ 1.840172e-01,  1.075276e-01,  6.186187e-02,  1.495719e-01, ...,  2.777254e-02, -3.337724e-02,\n",
       "         -5.033139e-02, -7.935667e-02],\n",
       "        [ 6.584296e-02,  8.974664e-02,  8.604165e-02,  9.996102e-02, ...,  2.057362e-01,  2.437396e-01,\n",
       "          2.530226e-01,  2.004630e-01],\n",
       "        [ 9.364848e-02,  2.017706e-01,  2.386461e-01,  1.540475e-01, ...,  2.210179e-01,  2.937905e-01,\n",
       "          3.533009e-01,  3.337061e-01]],\n",
       "\n",
       "       [[ 1.175666e+00,  1.220566e+00,  1.200156e+00,  1.138659e+00, ...,  2.616597e-01,  1.827672e-02,\n",
       "         -1.212664e-01, -2.315757e-01],\n",
       "        [ 1.205749e+00,  1.243035e+00,  1.218173e+00,  1.190289e+00, ...,  2.978287e-01,  3.892785e-02,\n",
       "         -6.948082e-02, -1.628896e-01],\n",
       "        [ 1.199235e+00,  1.229218e+00,  1.232467e+00,  1.257338e+00, ...,  3.268813e-01,  1.160583e-01,\n",
       "          4.168704e-03, -6.838660e-02],\n",
       "        [ 1.137234e+00,  1.183527e+00,  1.212620e+00,  1.240344e+00, ...,  3.642857e-01,  1.664781e-01,\n",
       "          6.743429e-02,  5.027696e-03],\n",
       "        ...,\n",
       "        [-6.436907e-01, -5.503951e-01, -4.742598e-01, -4.875618e-01, ..., -1.037436e+00, -1.040614e+00,\n",
       "         -1.177142e+00, -1.220379e+00],\n",
       "        [-7.436590e-01, -6.238915e-01, -4.735585e-01, -5.080217e-01, ..., -8.872295e-01, -7.897852e-01,\n",
       "         -9.088947e-01, -9.806531e-01],\n",
       "        [-6.747395e-01, -6.015898e-01, -4.995871e-01, -5.259740e-01, ..., -9.087682e-01, -8.189643e-01,\n",
       "         -9.191239e-01, -9.423386e-01],\n",
       "        [-5.969926e-01, -6.313481e-01, -6.231820e-01, -5.786868e-01, ..., -9.600114e-01, -9.989497e-01,\n",
       "         -1.157557e+00, -1.163065e+00]],\n",
       "\n",
       "       [[-6.557242e-01, -6.066410e-01, -6.021546e-01, -5.794507e-01, ..., -4.527237e-01, -4.429210e-01,\n",
       "         -4.480021e-01, -4.591919e-01],\n",
       "        [-6.415199e-01, -6.098990e-01, -5.912560e-01, -5.577773e-01, ..., -4.720502e-01, -4.568951e-01,\n",
       "         -4.618532e-01, -4.677817e-01],\n",
       "        [-6.225531e-01, -6.006424e-01, -5.806381e-01, -5.501240e-01, ..., -4.918783e-01, -4.865147e-01,\n",
       "         -4.856077e-01, -4.998766e-01],\n",
       "        [-6.047767e-01, -5.906608e-01, -5.885633e-01, -5.494895e-01, ..., -5.143435e-01, -5.171611e-01,\n",
       "         -5.178832e-01, -5.292745e-01],\n",
       "        ...,\n",
       "        [-1.869215e-02,  2.220136e-02, -2.010003e-02, -5.067983e-02, ...,  7.476941e-02,  7.811081e-02,\n",
       "          5.143912e-02,  3.505984e-02],\n",
       "        [ 4.171029e-03,  2.213848e-02, -1.570334e-02, -3.041671e-02, ...,  7.305312e-02,  8.682191e-02,\n",
       "          4.202220e-02,  1.745483e-02],\n",
       "        [ 7.760525e-05, -2.853160e-02,  2.213925e-03,  1.826693e-02, ..., -1.318067e-03,  2.810496e-02,\n",
       "         -3.546959e-02, -7.692659e-02],\n",
       "        [-3.465828e-02, -6.448114e-02,  1.741180e-02,  5.688281e-02, ...,  2.353421e-02,  5.509746e-02,\n",
       "         -2.259058e-02, -8.520116e-02]]], dtype=float32)"
      ]
     },
     "execution_count": 179,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf_output[0][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 180,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[ 5.928250e-01,  6.495560e-01,  6.586254e-01,  6.328704e-01, ...,  1.145797e-01, -1.446192e-02,\n",
       "         -9.115485e-02, -1.700499e-01],\n",
       "        [ 6.237740e-01,  6.704966e-01,  6.704947e-01,  6.700016e-01, ...,  1.083156e-01, -3.941596e-02,\n",
       "         -9.333248e-02, -1.599272e-01],\n",
       "        [ 6.543528e-01,  7.007562e-01,  6.936352e-01,  7.207102e-01, ...,  1.495145e-01, -5.462829e-04,\n",
       "         -5.724194e-02, -1.157975e-01],\n",
       "        [ 6.586937e-01,  6.734817e-01,  6.918388e-01,  7.479682e-01, ...,  1.761432e-01,  5.087249e-02,\n",
       "         -1.233188e-02, -3.808299e-02],\n",
       "        ...,\n",
       "        [-5.286742e-01, -3.994761e-01, -4.185633e-01, -5.430715e-01, ..., -8.238600e-01, -8.971226e-01,\n",
       "         -1.071087e+00, -1.090646e+00],\n",
       "        [-6.421031e-01, -5.654278e-01, -4.625924e-01, -5.476409e-01, ..., -8.586741e-01, -7.533782e-01,\n",
       "         -8.619232e-01, -9.086895e-01],\n",
       "        [-6.917681e-01, -6.695323e-01, -5.033025e-01, -4.304609e-01, ..., -8.144889e-01, -6.926728e-01,\n",
       "         -7.351253e-01, -7.078845e-01],\n",
       "        [-5.975550e-01, -6.254424e-01, -5.287780e-01, -4.007401e-01, ..., -7.700869e-01, -7.775310e-01,\n",
       "         -9.107289e-01, -8.765057e-01]],\n",
       "\n",
       "       [[ 9.304108e-01,  9.061006e-01,  8.904200e-01,  8.748452e-01, ...,  1.507193e-01,  4.863602e-02,\n",
       "         -2.965286e-02, -6.887357e-02],\n",
       "        [ 9.354134e-01,  9.211414e-01,  9.025218e-01,  8.819546e-01, ...,  1.583953e-01,  7.783134e-02,\n",
       "          2.384748e-03, -2.762413e-02],\n",
       "        [ 9.247492e-01,  9.207308e-01,  9.197693e-01,  9.268373e-01, ...,  1.964801e-01,  1.070963e-01,\n",
       "          6.585616e-02,  3.885203e-02],\n",
       "        [ 8.749162e-01,  9.000026e-01,  9.111722e-01,  9.365388e-01, ...,  2.133995e-01,  1.541420e-01,\n",
       "          1.271267e-01,  1.043288e-01],\n",
       "        ...,\n",
       "        [-4.185115e-01, -4.692028e-01, -4.844882e-01, -4.011375e-01, ..., -9.281288e-01, -9.043365e-01,\n",
       "         -8.730649e-01, -8.296068e-01],\n",
       "        [-4.699255e-01, -4.824458e-01, -4.375248e-01, -3.610857e-01, ..., -6.990985e-01, -7.424051e-01,\n",
       "         -7.382967e-01, -6.652063e-01],\n",
       "        [-5.109918e-01, -4.493734e-01, -4.005437e-01, -4.306577e-01, ..., -6.778324e-01, -7.774227e-01,\n",
       "         -8.163334e-01, -7.388382e-01],\n",
       "        [-4.955843e-01, -4.066808e-01, -4.318877e-01, -5.693138e-01, ..., -6.130958e-01, -7.152777e-01,\n",
       "         -7.512010e-01, -6.329308e-01]],\n",
       "\n",
       "       [[-8.835173e-01, -8.811719e-01, -8.571362e-01, -8.381184e-01, ..., -1.226337e-01, -1.413189e-02,\n",
       "          6.252305e-02,  1.186951e-01],\n",
       "        [-8.997300e-01, -9.013266e-01, -8.836781e-01, -8.679811e-01, ..., -1.311588e-01, -3.768335e-02,\n",
       "          2.365048e-02,  6.502268e-02],\n",
       "        [-8.955600e-01, -9.047340e-01, -9.121482e-01, -9.156053e-01, ..., -1.767565e-01, -9.044454e-02,\n",
       "         -3.904480e-02, -1.396201e-03],\n",
       "        [-8.667954e-01, -8.885559e-01, -9.212384e-01, -9.478627e-01, ..., -2.214962e-01, -1.539418e-01,\n",
       "         -1.061586e-01, -7.190461e-02],\n",
       "        ...,\n",
       "        [ 6.527845e-01,  6.506559e-01,  6.483269e-01,  6.090642e-01, ...,  1.127803e+00,  1.157625e+00,\n",
       "          1.176368e+00,  1.130930e+00],\n",
       "        [ 7.203699e-01,  6.666610e-01,  6.309564e-01,  6.287088e-01, ...,  9.141809e-01,  9.502071e-01,\n",
       "          9.911880e-01,  9.432814e-01],\n",
       "        [ 7.402143e-01,  6.619299e-01,  6.260362e-01,  6.843429e-01, ...,  9.217322e-01,  9.615628e-01,\n",
       "          1.000422e+00,  9.200149e-01],\n",
       "        [ 7.193744e-01,  6.757457e-01,  6.706685e-01,  7.347934e-01, ...,  9.194486e-01,  9.976853e-01,\n",
       "          1.031135e+00,  9.027116e-01]],\n",
       "\n",
       "       [[-5.611107e-01, -6.185374e-01, -6.274537e-01, -6.018535e-01, ..., -3.286290e-01, -1.958319e-01,\n",
       "         -1.233654e-01, -5.534839e-02],\n",
       "        [-5.882207e-01, -6.269063e-01, -6.223419e-01, -5.886477e-01, ..., -2.825735e-01, -1.406765e-01,\n",
       "         -6.240558e-02, -1.910090e-03],\n",
       "        [-6.440452e-01, -6.522714e-01, -6.412833e-01, -5.956016e-01, ..., -2.842835e-01, -1.286600e-01,\n",
       "         -5.839922e-02, -1.226380e-02],\n",
       "        [-6.646396e-01, -6.643090e-01, -6.534098e-01, -6.352038e-01, ..., -3.086022e-01, -1.643063e-01,\n",
       "         -1.076260e-01, -7.506974e-02],\n",
       "        ...,\n",
       "        [ 1.974958e-01,  1.334879e-01,  2.476262e-01,  3.490084e-01, ...,  9.260128e-01,  9.112363e-01,\n",
       "          8.825827e-01,  8.980755e-01],\n",
       "        [ 2.174751e-01,  2.777752e-01,  3.892178e-01,  3.058348e-01, ...,  6.894872e-01,  6.306827e-01,\n",
       "          6.324791e-01,  7.319039e-01],\n",
       "        [ 4.317094e-01,  5.105460e-01,  4.332957e-01,  2.208417e-01, ...,  6.646112e-01,  5.065519e-01,\n",
       "          4.617480e-01,  6.012140e-01],\n",
       "        [ 4.610324e-01,  4.142273e-01,  2.067479e-01,  1.124400e-01, ...,  3.580654e-01,  2.135994e-01,\n",
       "          1.807216e-01,  3.309866e-01]],\n",
       "\n",
       "       ...,\n",
       "\n",
       "       [[-1.021994e+00, -1.007560e+00, -1.020993e+00, -1.028300e+00, ..., -6.495625e-01, -6.656052e-01,\n",
       "         -6.676931e-01, -6.754252e-01],\n",
       "        [-1.022150e+00, -1.020793e+00, -1.031548e+00, -1.039426e+00, ..., -6.874483e-01, -6.935481e-01,\n",
       "         -7.095402e-01, -7.143820e-01],\n",
       "        [-1.000767e+00, -1.022446e+00, -1.045676e+00, -1.072399e+00, ..., -7.137644e-01, -7.332122e-01,\n",
       "         -7.495062e-01, -7.635692e-01],\n",
       "        [-9.922659e-01, -1.020445e+00, -1.055791e+00, -1.083923e+00, ..., -7.817140e-01, -7.956998e-01,\n",
       "         -8.141065e-01, -8.214797e-01],\n",
       "        ...,\n",
       "        [ 5.825371e-01,  5.784728e-01,  5.523584e-01,  5.299688e-01, ...,  1.002575e+00,  1.063073e+00,\n",
       "          1.038671e+00,  9.423774e-01],\n",
       "        [ 5.757274e-01,  5.073735e-01,  4.978513e-01,  5.697755e-01, ...,  7.133898e-01,  8.178766e-01,\n",
       "          8.427096e-01,  7.569195e-01],\n",
       "        [ 6.112009e-01,  5.420885e-01,  5.843502e-01,  6.813412e-01, ...,  8.724557e-01,  9.414760e-01,\n",
       "          9.315102e-01,  8.196849e-01],\n",
       "        [ 6.394557e-01,  6.162713e-01,  6.582525e-01,  6.849660e-01, ...,  7.615169e-01,  8.150862e-01,\n",
       "          7.435393e-01,  5.550404e-01]],\n",
       "\n",
       "       [[ 4.288197e-02,  3.586289e-02,  1.979762e-02,  4.157487e-02, ...,  1.990840e-01,  2.420095e-01,\n",
       "          2.877120e-01,  3.103569e-01],\n",
       "        [ 4.962894e-02,  3.502490e-02,  3.614373e-03,  7.033223e-05, ...,  1.563922e-01,  1.922820e-01,\n",
       "          2.452845e-01,  2.568748e-01],\n",
       "        [ 7.724187e-02,  4.813374e-02,  1.849236e-02, -1.296011e-02, ...,  1.404089e-01,  1.845620e-01,\n",
       "          2.243527e-01,  2.408762e-01],\n",
       "        [ 1.300122e-01,  1.025737e-01,  8.795914e-02,  4.576286e-02, ...,  1.721860e-01,  2.096085e-01,\n",
       "          2.493705e-01,  2.706239e-01],\n",
       "        ...,\n",
       "        [ 2.466994e-01,  1.628171e-01,  1.162946e-01,  1.563058e-01, ..., -6.753094e-02, -9.556959e-02,\n",
       "         -4.442038e-02, -9.396417e-03],\n",
       "        [ 1.840173e-01,  1.075276e-01,  6.186185e-02,  1.495719e-01, ...,  2.777258e-02, -3.337731e-02,\n",
       "         -5.033160e-02, -7.935677e-02],\n",
       "        [ 6.584295e-02,  8.974665e-02,  8.604161e-02,  9.996101e-02, ...,  2.057362e-01,  2.437397e-01,\n",
       "          2.530226e-01,  2.004631e-01],\n",
       "        [ 9.364849e-02,  2.017706e-01,  2.386460e-01,  1.540475e-01, ...,  2.210179e-01,  2.937904e-01,\n",
       "          3.533009e-01,  3.337062e-01]],\n",
       "\n",
       "       [[ 1.175666e+00,  1.220566e+00,  1.200156e+00,  1.138659e+00, ...,  2.616597e-01,  1.827675e-02,\n",
       "         -1.212664e-01, -2.315756e-01],\n",
       "        [ 1.205749e+00,  1.243035e+00,  1.218173e+00,  1.190289e+00, ...,  2.978287e-01,  3.892778e-02,\n",
       "         -6.948082e-02, -1.628896e-01],\n",
       "        [ 1.199235e+00,  1.229218e+00,  1.232467e+00,  1.257338e+00, ...,  3.268813e-01,  1.160583e-01,\n",
       "          4.168683e-03, -6.838661e-02],\n",
       "        [ 1.137234e+00,  1.183527e+00,  1.212620e+00,  1.240344e+00, ...,  3.642857e-01,  1.664781e-01,\n",
       "          6.743427e-02,  5.027659e-03],\n",
       "        ...,\n",
       "        [-6.436908e-01, -5.503952e-01, -4.742598e-01, -4.875617e-01, ..., -1.037436e+00, -1.040613e+00,\n",
       "         -1.177142e+00, -1.220379e+00],\n",
       "        [-7.436591e-01, -6.238915e-01, -4.735585e-01, -5.080217e-01, ..., -8.872296e-01, -7.897854e-01,\n",
       "         -9.088947e-01, -9.806532e-01],\n",
       "        [-6.747395e-01, -6.015897e-01, -4.995870e-01, -5.259740e-01, ..., -9.087684e-01, -8.189642e-01,\n",
       "         -9.191239e-01, -9.423386e-01],\n",
       "        [-5.969926e-01, -6.313480e-01, -6.231819e-01, -5.786868e-01, ..., -9.600114e-01, -9.989499e-01,\n",
       "         -1.157557e+00, -1.163065e+00]],\n",
       "\n",
       "       [[-6.557240e-01, -6.066412e-01, -6.021547e-01, -5.794508e-01, ..., -4.527237e-01, -4.429210e-01,\n",
       "         -4.480022e-01, -4.591919e-01],\n",
       "        [-6.415200e-01, -6.098990e-01, -5.912560e-01, -5.577772e-01, ..., -4.720501e-01, -4.568951e-01,\n",
       "         -4.618531e-01, -4.677817e-01],\n",
       "        [-6.225531e-01, -6.006425e-01, -5.806381e-01, -5.501241e-01, ..., -4.918782e-01, -4.865148e-01,\n",
       "         -4.856077e-01, -4.998765e-01],\n",
       "        [-6.047767e-01, -5.906608e-01, -5.885634e-01, -5.494895e-01, ..., -5.143436e-01, -5.171611e-01,\n",
       "         -5.178832e-01, -5.292746e-01],\n",
       "        ...,\n",
       "        [-1.869220e-02,  2.220134e-02, -2.010003e-02, -5.067986e-02, ...,  7.476944e-02,  7.811080e-02,\n",
       "          5.143915e-02,  3.505992e-02],\n",
       "        [ 4.170997e-03,  2.213846e-02, -1.570337e-02, -3.041676e-02, ...,  7.305323e-02,  8.682193e-02,\n",
       "          4.202233e-02,  1.745484e-02],\n",
       "        [ 7.753043e-05, -2.853164e-02,  2.213888e-03,  1.826694e-02, ..., -1.318107e-03,  2.810498e-02,\n",
       "         -3.546954e-02, -7.692661e-02],\n",
       "        [-3.465827e-02, -6.448118e-02,  1.741178e-02,  5.688276e-02, ...,  2.353423e-02,  5.509750e-02,\n",
       "         -2.259056e-02, -8.520117e-02]]], dtype=float32)"
      ]
     },
     "execution_count": 180,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch_output[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 181,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# torch with TF\n",
    "if not inputs_same(torch_output[0], tf_output[0][0]):\n",
    "    raise Exception(\"PyTorch and TF outputs are not same\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 182,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf_rep.export_graph(MODELS/f\"{arch_name}.pb\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 183,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('0', '4')"
      ]
     },
     "execution_count": 183,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf_rep.inputs[0], tf_rep.outputs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 184,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(<tf.Tensor '0:0' shape=(1, 224, 224, 3) dtype=float32>,\n",
       " <tf.Tensor 'transpose_3:0' shape=(1, 64, 222, 222) dtype=float32>)"
      ]
     },
     "execution_count": 184,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf_rep.tensor_dict[tf_rep.inputs[0]], tf_rep.tensor_dict[ tf_rep.outputs[0]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ONNX caffe2 backend"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 185,
   "metadata": {},
   "outputs": [],
   "source": [
    "import caffe2.python.onnx.backend as onnx_caffe2_backend\n",
    "\n",
    "# content image shape 224, 224, 3\n",
    "model = onnx.load(MODELS/f'{arch_name}.onnx')\n",
    "prepared_backend = onnx_caffe2_backend.prepare(model, device='CUDA' if torch.cuda.is_available() else 'CPU')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 186,
   "metadata": {},
   "outputs": [],
   "source": [
    "inp = {model.graph.input[0].name: numpy_input.astype(np.float32)}\n",
    "caffe2_output = prepared_backend.run(inp)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 187,
   "metadata": {},
   "outputs": [],
   "source": [
    "# torch with TF\n",
    "if not inputs_same(torch_output[0], caffe2_output[0]):\n",
    "    raise Exception(\"PyTorch and Caffe2 outputs are not same\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ONNX to caffe2 pb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 188,
   "metadata": {},
   "outputs": [],
   "source": [
    "# extract the workspace and the model proto from the internal representation\n",
    "c2_workspace = prepared_backend.workspace\n",
    "c2_model = prepared_backend.predict_net\n",
    "\n",
    "# Now import the caffe2 mobile exporter\n",
    "from caffe2.python.predictor import mobile_exporter\n",
    "\n",
    "# call the Export to get the predict_net, init_net. These nets are needed for running things on mobile\n",
    "init_net, predict_net = mobile_exporter.Export(c2_workspace, c2_model, c2_model.external_input)\n",
    "\n",
    "# Let's also save the init_net and predict_net to a file that we will later use for running them on mobile\n",
    "with open(MODELS/f'{arch_name}_init_net.pb', \"wb\") as fopen:\n",
    "    fopen.write(init_net.SerializeToString())\n",
    "with open(MODELS/f'{arch_name}_net.pb', \"wb\") as fopen:\n",
    "    fopen.write(predict_net.SerializeToString())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 189,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch-jit-export_predict = core.Net('torch-jit-export_predict')\n",
      "torch-jit-export_predict.Transpose(['0'], ['3'], axes=[0, 3, 1, 2])\n",
      "torch-jit-export_predict.Conv(['3', '1', '2'], ['4'], strides=[1, 1], pads=[0, 0, 0, 0], kernels=[3, 3], group=1, dilations=[1, 1])\n"
     ]
    }
   ],
   "source": [
    "# Some standard imports\n",
    "from caffe2.proto import caffe2_pb2\n",
    "from caffe2.python import core, net_drawer, net_printer, visualize, workspace, utils\n",
    "\n",
    "# Let's run the mobile nets that we generated above so that caffe2 workspace is properly initialized\n",
    "workspace.RunNetOnce(init_net)\n",
    "workspace.RunNetOnce(predict_net)\n",
    "\n",
    "# Caffe2 has a nice net_printer to be able to inspect what the net looks like and identify\n",
    "# what our input and output blob names are.\n",
    "print(net_printer.to_string(predict_net))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 190,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now, let's also pass in the resized cat image for processing by the model.\n",
    "workspace.FeedBlob(tf_rep.inputs[0], numpy_input.astype(np.float32))\n",
    "\n",
    "# run the predict_net to get the model output\n",
    "workspace.RunNetOnce(predict_net)\n",
    "\n",
    "# Now let's get the model output blob\n",
    "caffe2_proto_output = workspace.FetchBlob(tf_rep.outputs[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 191,
   "metadata": {},
   "outputs": [],
   "source": [
    "# torch with TF\n",
    "if not inputs_same(torch_output[0], caffe2_proto_output[0]):\n",
    "    raise Exception(\"PyTorch and Caffe2 Proto outputs are not same\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TF to TFlite\n",
    "\n",
    "**Check failed: status.ok() Unexpected value for attribute 'data_format'. \n",
    "Expected 'NHWC'**\n",
    "\n",
    "channels first (NCHW) is normally optimal on GPU and channels last (NHWC)\n",
    "on CPU. The exception is Intel MKL on CPU which is optimal with\n",
    "channels_last.\n",
    "\n",
    "At the moment TensorFlow Lite supports only TensorFlow's \"NHWC\" format, and broadcasting is only support in a limited number of ops (tf.add, tf.mul, tf.sub, and tf.div).\n",
    "\n",
    "https://www.tensorflow.org/lite/guide/ops_compatibility\n",
    "\n",
    "\n",
    "When input is \"NHWC\" there is a different error:\n",
    "\n",
    "- 2019-03-19 18:24:24.038521: I tensorflow/lite/toco/import_tensorflow.cc:1324] Converting unsupported operation: PyFunc\n",
    "\n",
    "- 2019-03-19 18:24:24.051107: F tensorflow/lite/toco/import_tensorflow.cc:114] Check failed: attr.value_case() == AttrValue::kType (1 vs. 6)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 192,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9452"
      ]
     },
     "execution_count": 192,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "converter = tf.lite.TFLiteConverter.from_frozen_graph(str(MODELS/f'{arch_name}.pb'),\n",
    "                                                      tf_rep.tensor_dict[tf_rep.inputs[0]].name.split(\":\")[:1],\n",
    "                                                      tf_rep.tensor_dict[tf_rep.outputs[0]].name.split(\":\")[:1])\n",
    "tflite_model = converter.convert()\n",
    "open(MODELS/f\"{arch_name}.tflite\", \"wb\").write(tflite_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}