tutorials/CorrectnessVerificationAndPerformanceComparison.ipynb

{
 "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": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:This caffe2 python run does not have GPU support. Will run in CPU only mode.\n",
      "WARNING:root:Debug message: No module named caffe2_pybind11_state_gpu\n"
     ]
    }
   ],
   "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 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, \\\n",
    "    benchmark_caffe2_model, 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(%0 : Float(3, 1, 28, 28)\n",
      "      %1 : Float(10, 1, 5, 5)\n",
      "      %2 : Float(10)\n",
      "      %3 : Float(20, 10, 5, 5)\n",
      "      %4 : Float(20)\n",
      "      %5 : Float(50, 320)\n",
      "      %6 : Float(50)\n",
      "      %7 : Float(10, 50)\n",
      "      %8 : Float(10)) {\n",
      "  %9 : UNKNOWN_TYPE = Conv[kernel_shape=[5, 5], strides=[1, 1], pads=[0, 0, 0, 0], dilations=[1, 1], group=1](%0, %1)\n",
      "  %10 : Float(3, 10, 24, 24) = Add[broadcast=1, axis=1](%9, %2)\n",
      "  %11 : Float(3, 10, 12, 12) = MaxPool[kernel_shape=[2, 2], pads=[0, 0], strides=[2, 2]](%10)\n",
      "  %12 : Float(3, 10, 12, 12) = Relu(%11)\n",
      "  %13 : UNKNOWN_TYPE = Conv[kernel_shape=[5, 5], strides=[1, 1], pads=[0, 0, 0, 0], dilations=[1, 1], group=1](%12, %3)\n",
      "  %14 : Float(3, 20, 8, 8) = Add[broadcast=1, axis=1](%13, %4)\n",
      "  %15 : Float(3, 20, 4, 4) = MaxPool[kernel_shape=[2, 2], pads=[0, 0], strides=[2, 2]](%14)\n",
      "  %16 : Float(3, 20, 4, 4) = Relu(%15)\n",
      "  %17 : Float(3, 320) = Reshape[shape=[-1, 320]](%16)\n",
      "  %18 : Float(320!, 50!) = Transpose[perm=[1, 0]](%5)\n",
      "  %20 : Float(3, 50) = Gemm[alpha=1, beta=1, broadcast=1](%17, %18, %6)\n",
      "  %21 : Float(3, 50) = Relu(%20)\n",
      "  %22 : Float(3, 50), %23 : UNKNOWN_TYPE = Dropout[is_test=1, ratio=0.5](%21)\n",
      "  %24 : Float(50!, 10!) = Transpose[perm=[1, 0]](%7)\n",
      "  %26 : Float(3, 10) = Gemm[alpha=1, beta=1, broadcast=1](%22, %24, %8)\n",
      "  %27 : Float(3, 10) = Softmax[axis=1](%26)\n",
      "  %28 : Float(3, 10) = Log(%27)\n",
      "  return (%28);\n",
      "}\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.graph, 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": {
    "collapsed": true
   },
   "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": {
    "collapsed": true
   },
   "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": "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": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PyTorch model's execution time is 0.580716133118 milliseconds/ iteration, 1722.011742 iterations per second.\n",
      "Caffe2 model's execution time is 0.399529695511 milliseconds / iteration, 2502.94286316 iterations per second\n"
     ]
    }
   ],
   "source": [
    "pytorch_time = benchmark_pytorch_model(pytorch_model, inputs)\n",
    "caffe2_time = benchmark_caffe2_model(init_net, predict_net)\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))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}