{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Verify the Correctness of Exported Model and Compare the Performance\n",
    "\n",
    "In this tutorial, we are going to show:\n",
    "- how to verify the correctness of the exported model\n",
    "- how to compare the performance with the original model\n",
    "\n",
    "We choose PyTorch to export the ONNX model, and use Caffe2 as the backend.\n",
    "After that, the outputs and performance of two models are compared.\n",
    "\n",
    "To run this tutorial, please make sure that `caffe2`, `pytorch` and `onnx` are already installed.\n",
    "\n",
    "First, let's create a PyTorch model and prepare the inputs of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import absolute_import\n",
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "from __future__ import unicode_literals\n",
    "\n",
    "import io\n",
    "import time\n",
    "import numpy as np\n",
    "import torch\n",
    "import onnx\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from caffe2.proto import caffe2_pb2\n",
    "from caffe2.python import core\n",
    "from torch.autograd import Variable\n",
    "from caffe2.python.onnx.backend import Caffe2Backend\n",
    "from caffe2.python.onnx.helper import c2_native_run_net, save_caffe2_net, load_caffe2_net,benchmark_pytorch_model\n",
    "\n",
    "\n",
    "class MNIST(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super(MNIST, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)\n",
    "        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)\n",
    "        self.conv2_drop = nn.Dropout2d()\n",
    "        self.fc1 = nn.Linear(320, 50)\n",
    "        self.fc2 = nn.Linear(50, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.relu(F.max_pool2d(self.conv1(x), 2))\n",
    "        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))\n",
    "        x = x.view(-1, 320)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.dropout(x, training=self.training)\n",
    "        x = self.fc2(x)\n",
    "        return F.log_softmax(x)\n",
    "\n",
    "\n",
    "# Create a pytorch model.\n",
    "pytorch_model = MNIST()\n",
    "pytorch_model.train(False)\n",
    "\n",
    "# Make the inputs in tuple format.\n",
    "inputs = (Variable(torch.randn(3, 1, 28, 28), requires_grad=True), )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the PyTorch exporter to generate an ONNX model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "graph(%input.1 : Float(3, 1, 28, 28),\n",
      "      %conv1.weight : Float(10, 1, 5, 5),\n",
      "      %conv1.bias : Float(10),\n",
      "      %conv2.weight : Float(20, 10, 5, 5),\n",
      "      %conv2.bias : Float(20),\n",
      "      %fc1.weight : Float(50, 320),\n",
      "      %fc1.bias : Float(50),\n",
      "      %fc2.weight : Float(10, 50),\n",
      "      %fc2.bias : Float(10)):\n",
      "  %9 : Float(3, 10, 24, 24) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[5, 5], pads=[0, 0, 0, 0], strides=[1, 1]](%input.1, %conv1.weight, %conv1.bias), scope: MNIST/Conv2d[conv1] # /home/marouenez/anaconda3/envs/masterthesis/lib/python3.7/site-packages/torch/nn/modules/conv.py:340:0\n",
      "  %10 : Float(3, 10, 12, 12) = onnx::MaxPool[kernel_shape=[2, 2], pads=[0, 0, 0, 0], strides=[2, 2]](%9), scope: MNIST # /home/marouenez/anaconda3/envs/masterthesis/lib/python3.7/site-packages/torch/nn/functional.py:487:0\n",
      "  %11 : Float(3, 10, 12, 12) = onnx::Relu(%10), scope: MNIST # /home/marouenez/anaconda3/envs/masterthesis/lib/python3.7/site-packages/torch/nn/functional.py:913:0\n",
      "  %12 : Float(3, 20, 8, 8) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[5, 5], pads=[0, 0, 0, 0], strides=[1, 1]](%11, %conv2.weight, %conv2.bias), scope: MNIST/Dropout2d[conv2_drop] # /home/marouenez/anaconda3/envs/masterthesis/lib/python3.7/site-packages/torch/nn/functional.py:844:0\n",
      "  %13 : Float(3, 20, 4, 4) = onnx::MaxPool[kernel_shape=[2, 2], pads=[0, 0, 0, 0], strides=[2, 2]](%12), scope: MNIST # /home/marouenez/anaconda3/envs/masterthesis/lib/python3.7/site-packages/torch/nn/functional.py:487:0\n",
      "  %14 : Float(3, 20, 4, 4) = onnx::Relu(%13), scope: MNIST # /home/marouenez/anaconda3/envs/masterthesis/lib/python3.7/site-packages/torch/nn/functional.py:913:0\n",
      "  %15 : Tensor = onnx::Constant[value=  -1  320 [ Variable[CPULongType]{2} ]](), scope: MNIST\n",
      "  %16 : Float(3, 320) = onnx::Reshape(%14, %15), scope: MNIST # <ipython-input-1-c1a2123bd117>:34:0\n",
      "  %17 : Float(3, 50) = onnx::Gemm[alpha=1, beta=1, transB=1](%16, %fc1.weight, %fc1.bias), scope: MNIST/Linear[fc1] # /home/marouenez/anaconda3/envs/masterthesis/lib/python3.7/site-packages/torch/nn/functional.py:1369:0\n",
      "  %18 : Float(3, 50) = onnx::Relu(%17), scope: MNIST # /home/marouenez/anaconda3/envs/masterthesis/lib/python3.7/site-packages/torch/nn/functional.py:806:0\n",
      "  %19 : Float(3, 10) = onnx::Gemm[alpha=1, beta=1, transB=1](%18, %fc2.weight, %fc2.bias), scope: MNIST/Linear[fc2] # /home/marouenez/anaconda3/envs/masterthesis/lib/python3.7/site-packages/torch/nn/functional.py:1369:0\n",
      "  %20 : Float(3, 10) = onnx::LogSoftmax[axis=1](%19), scope: MNIST # /home/marouenez/anaconda3/envs/masterthesis/lib/python3.7/site-packages/torch/nn/functional.py:1316:0\n",
      "  return (%20)\n",
      "\n",
      "Check the ONNX model.\n"
     ]
    }
   ],
   "source": [
    "# Export an ONNX model.\n",
    "f = io.BytesIO()\n",
    "torch.onnx.export(pytorch_model, inputs, f, verbose=True)\n",
    "onnx_model = onnx.ModelProto.FromString(f.getvalue())\n",
    "\n",
    "# Check whether the onnx_model is valid or not.\n",
    "print(\"Check the ONNX model.\")\n",
    "onnx.checker.check_model(onnx_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we have an ONNX model, let's turn it into a Caffe2 one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Convert the model to a Caffe2 model.\n"
     ]
    }
   ],
   "source": [
    "# Convert the ONNX model to a Caffe2 model.\n",
    "print(\"Convert the model to a Caffe2 model.\")\n",
    "init_net, predict_net = Caffe2Backend.onnx_graph_to_caffe2_net(onnx_model, device=\"CPU\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Caffe2 takes a list of numpy array as inputs. So we need to change the format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepare the inputs for Caffe2.\n",
    "caffe2_inputs = [var.data.numpy() for var in inputs]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code shows how to save and load a Caffe2 model. It is purely for demonstration purpose here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the converted Caffe2 model in the protobuf files. (Optional)\n",
    "init_file = \"./output/mymodel_init.pb\"\n",
    "predict_file = \"./output/mymodel_predict.pb\"\n",
    "save_caffe2_net(init_net, init_file, output_txt=False)\n",
    "save_caffe2_net(predict_net, predict_file, output_txt=True)\n",
    "\n",
    "# Load the Caffe2 model.\n",
    "init_net = load_caffe2_net(init_file)\n",
    "predict_net = load_caffe2_net(predict_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run PyTorch and Caffe2 models separately, and get the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Compute the results using the PyTorch model.\n",
    "pytorch_results = pytorch_model(*inputs)\n",
    "\n",
    "# Compute the results using the Caffe2 model.\n",
    "_, caffe2_results = c2_native_run_net(init_net, predict_net, caffe2_inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have the results, let's check the correctness of the exported model.\n",
    "If no assertion fails, our model has achieved expected precision."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The exported model achieves 5-decimal precision.\n"
     ]
    }
   ],
   "source": [
    "# Check the decimal precision of the exported Caffe2.\n",
    "expected_decimal = 5\n",
    "for p, c in zip([pytorch_results], caffe2_results):\n",
    "    np.testing.assert_almost_equal(p.data.cpu().numpy(), c, decimal=expected_decimal)\n",
    "print(\"The exported model achieves {}-decimal precision.\".format(expected_decimal))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def benchmark_caffe2_model(init_net, predict_net,inputs, warmup_iters=3, main_iters=10):\n",
    "    '''\n",
    "     Run the model several times, and measure the execution time.\n",
    "     Print the execution time per iteration (millisecond) and the number of iterations per second.\n",
    "    '''\n",
    "    for _i in range(warmup_iters):\n",
    "        ws, caffe2_results = c2_native_run_net(init_net, predict_net, inputs)    \n",
    "\n",
    "    total_time = 0.0\n",
    "    for _i in range(main_iters):\n",
    "        ts = time.time()\n",
    "        ws, caffe2_results = c2_native_run_net(init_net, predict_net, inputs)\n",
    "        te = time.time()\n",
    "        total_time += te - ts\n",
    "        \n",
    "    return total_time / main_iters * 1000"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code measures the performance of PyTorch and Caffe2 models.\n",
    "We report:\n",
    "- Execution time per iteration\n",
    "- Iterations per second"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PyTorch model's execution time is 0.6218433380126953 milliseconds/ iteration, 1608.1220765278736 iterations per second.\n",
      "Caffe2 model's execution time is 3.189969062805176 milliseconds / iteration, 313.48266403581545 iterations per second\n"
     ]
    }
   ],
   "source": [
    "pytorch_time = benchmark_pytorch_model(pytorch_model, inputs)\n",
    "caffe2_time = benchmark_caffe2_model(init_net, predict_net, caffe2_inputs)\n",
    "\n",
    "print(\"PyTorch model's execution time is {} milliseconds/ iteration, {} iterations per second.\".format(\n",
    "    pytorch_time, 1000 / pytorch_time))\n",
    "print(\"Caffe2 model's execution time is {} milliseconds / iteration, {} iterations per second\".format(\n",
    "    caffe2_time, 1000 / caffe2_time))"
   ]
  },
  {
   "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}