{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# For tips on running notebooks in Google Colab, see\n# https://pytorch.org/tutorials/beginner/colab\n%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "[Introduction to ONNX](intro_onnx.html) \\|\\| **Exporting a PyTorch model\nto ONNX** \\|\\| [Extending the ONNX exporter operator\nsupport](onnx_registry_tutorial.html) \\|\\| [Export a model with control\nflow to ONNX](export_control_flow_model_to_onnx_tutorial.html)\n\nExport a PyTorch model to ONNX\n==============================\n\n**Author**: [Ti-Tai Wang](https://github.com/titaiwangms), [Justin\nChu](justinchu@microsoft.com), [Thiago\nCrepaldi](https://github.com/thiagocrepaldi).\n\n```{=html}\n<div style=\"background-color: #54c7ec; color: #fff; font-weight: 700; padding-left: 10px; padding-top: 5px; padding-bottom: 5px\"><strong>NOTE:</strong></div>\n```\n```{=html}\n<div style=\"background-color: #f3f4f7; padding-left: 10px; padding-top: 10px; padding-bottom: 10px; padding-right: 10px\">\n```\n```{=html}\n<p>As of PyTorch 2.5, there are two versions of ONNX Exporter.</p>\n```\n```{=html}\n</div>\n```\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "In the [60 Minute\nBlitz](https://pytorch.org/tutorials/beginner/deep_learning_60min_blitz.html),\nwe had the opportunity to learn about PyTorch at a high level and train\na small neural network to classify images. In this tutorial, we are\ngoing to expand this to describe how to convert a model defined in\nPyTorch into the ONNX format using the\n`torch.onnx.export(..., dynamo=True)` ONNX exporter.\n\nWhile PyTorch is great for iterating on the development of models, the\nmodel can be deployed to production using different formats, including\n[ONNX](https://onnx.ai/) (Open Neural Network Exchange)!\n\nONNX is a flexible open standard format for representing machine\nlearning models which standardized representations of machine learning\nallow them to be executed across a gamut of hardware platforms and\nruntime environments from large-scale cloud-based supercomputers to\nresource-constrained edge devices, such as your web browser and phone.\n\nIn this tutorial, we'll learn how to:\n\n1.  Install the required dependencies.\n2.  Author a simple image classifier model.\n3.  Export the model to ONNX format.\n4.  Save the ONNX model in a file.\n5.  Visualize the ONNX model graph using\n    [Netron](https://github.com/lutzroeder/netron).\n6.  Execute the ONNX model with [ONNX Runtime]{.title-ref}\n7.  Compare the PyTorch results with the ones from the ONNX Runtime.\n\n1. Install the required dependencies\n====================================\n\nBecause the ONNX exporter uses `onnx` and `onnxscript` to translate\nPyTorch operators into ONNX operators, we will need to install them.\n\n> ``` {.bash}\n> pip install --upgrade onnx onnxscript\n> ```\n\n2. Author a simple image classifier model\n=========================================\n\nOnce your environment is set up, let's start modeling our image\nclassifier with PyTorch, exactly like we did in the [60 Minute\nBlitz](https://pytorch.org/tutorials/beginner/deep_learning_60min_blitz.html).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\nclass ImageClassifierModel(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.conv1 = nn.Conv2d(1, 6, 5)\n        self.conv2 = nn.Conv2d(6, 16, 5)\n        self.fc1 = nn.Linear(16 * 5 * 5, 120)\n        self.fc2 = nn.Linear(120, 84)\n        self.fc3 = nn.Linear(84, 10)\n\n    def forward(self, x: torch.Tensor):\n        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))\n        x = F.max_pool2d(F.relu(self.conv2(x)), 2)\n        x = torch.flatten(x, 1)\n        x = F.relu(self.fc1(x))\n        x = F.relu(self.fc2(x))\n        x = self.fc3(x)\n        return x"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "3. Export the model to ONNX format\n==================================\n\nNow that we have our model defined, we need to instantiate it and create\na random 32x32 input. Next, we can export the model to ONNX format.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "torch_model = ImageClassifierModel()\n# Create example inputs for exporting the model. The inputs should be a tuple of tensors.\nexample_inputs = (torch.randn(1, 1, 32, 32),)\nonnx_program = torch.onnx.export(torch_model, example_inputs, dynamo=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "3.5. (Optional) Optimize the ONNX model\n=======================================\n\nThe ONNX model can be optimized with constant folding, and elimination\nof redundant nodes. The optimization is done in-place, so the original\nONNX model is modified.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "onnx_program.optimize()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "As we can see, we didn\\'t need any code change to the model. The\nresulting ONNX model is stored within `torch.onnx.ONNXProgram` as a\nbinary protobuf file.\n\n4. Save the ONNX model in a file\n================================\n\nAlthough having the exported model loaded in memory is useful in many\napplications, we can save it to disk with the following code:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "onnx_program.save(\"image_classifier_model.onnx\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "You can load the ONNX file back into memory and check if it is well\nformed with the following code:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import onnx\n\nonnx_model = onnx.load(\"image_classifier_model.onnx\")\nonnx.checker.check_model(onnx_model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "5. Visualize the ONNX model graph using Netron\n==============================================\n\nNow that we have our model saved in a file, we can visualize it with\n[Netron](https://github.com/lutzroeder/netron). Netron can either be\ninstalled on macos, Linux or Windows computers, or run directly from the\nbrowser. Let\\'s try the web version by opening the following link:\n<https://netron.app/>.\n\n![image](https://pytorch.org/tutorials/_static/img/onnx/netron_web_ui.png){.align-center\nwidth=\"70.0%\"}\n\nOnce Netron is open, we can drag and drop our\n`image_classifier_model.onnx` file into the browser or select it after\nclicking the **Open model** button.\n\n![image](https://pytorch.org/tutorials/_static/img/onnx/image_classifier_onnx_model_on_netron_web_ui.png){width=\"50.0%\"}\n\nAnd that is it! We have successfully exported our PyTorch model to ONNX\nformat and visualized it with Netron.\n\n6. Execute the ONNX model with ONNX Runtime\n===========================================\n\nThe last step is executing the ONNX model with [ONNX\nRuntime]{.title-ref}, but before we do that, let\\'s install ONNX\nRuntime.\n\n> ``` {.bash}\n> pip install onnxruntime\n> ```\n\nThe ONNX standard does not support all the data structure and types that\nPyTorch does, so we need to adapt PyTorch input\\'s to ONNX format before\nfeeding it to ONNX Runtime. In our example, the input happens to be the\nsame, but it might have more inputs than the original PyTorch model in\nmore complex models.\n\nONNX Runtime requires an additional step that involves converting all\nPyTorch tensors to Numpy (in CPU) and wrap them on a dictionary with\nkeys being a string with the input name as key and the numpy tensor as\nthe value.\n\nNow we can create an *ONNX Runtime Inference Session*, execute the ONNX\nmodel with the processed input and get the output. In this tutorial,\nONNX Runtime is executed on CPU, but it could be executed on GPU as\nwell.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import onnxruntime\n\nonnx_inputs = [tensor.numpy(force=True) for tensor in example_inputs]\nprint(f\"Input length: {len(onnx_inputs)}\")\nprint(f\"Sample input: {onnx_inputs}\")\n\nort_session = onnxruntime.InferenceSession(\n    \"./image_classifier_model.onnx\", providers=[\"CPUExecutionProvider\"]\n)\n\nonnxruntime_input = {input_arg.name: input_value for input_arg, input_value in zip(ort_session.get_inputs(), onnx_inputs)}\n\n# ONNX Runtime returns a list of outputs\nonnxruntime_outputs = ort_session.run(None, onnxruntime_input)[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "7. Compare the PyTorch results with the ones from the ONNX Runtime\n==================================================================\n\nThe best way to determine whether the exported model is looking good is\nthrough numerical evaluation against PyTorch, which is our source of\ntruth.\n\nFor that, we need to execute the PyTorch model with the same input and\ncompare the results with ONNX Runtime\\'s. Before comparing the results,\nwe need to convert the PyTorch\\'s output to match ONNX\\'s format.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "torch_outputs = torch_model(*example_inputs)\n\nassert len(torch_outputs) == len(onnxruntime_outputs)\nfor torch_output, onnxruntime_output in zip(torch_outputs, onnxruntime_outputs):\n    torch.testing.assert_close(torch_output, torch.tensor(onnxruntime_output))\n\nprint(\"PyTorch and ONNX Runtime output matched!\")\nprint(f\"Output length: {len(onnxruntime_outputs)}\")\nprint(f\"Sample output: {onnxruntime_outputs}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Conclusion\n==========\n\nThat is about it! We have successfully exported our PyTorch model to\nONNX format, saved the model to disk, viewed it using Netron, executed\nit with ONNX Runtime and finally compared its numerical results with\nPyTorch\\'s.\n\nFurther reading\n===============\n\nThe list below refers to tutorials that ranges from basic examples to\nadvanced scenarios, not necessarily in the order they are listed. Feel\nfree to jump directly to specific topics of your interest or sit tight\nand have fun going through all of them to learn all there is about the\nONNX exporter.\n\n::: {.toctree hidden=\"\"}\n:::\n"
      ]
    }
  ],
  "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.10.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}