Automated tutorials push

Author: pytorchbot

_downloads/0bd6b9a8e47e1d64e4d20ef356a6095d/onnx_registry_tutorial.ipynb

@@ -0,0 +1,340 @@
+{
+ "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\n",
+    "to ONNX](export_simple_model_to_onnx_tutorial.html) \\|\\| [Extending the\n",
+    "ONNX exporter operator support](onnx_registry_tutorial.html) \\|\\|\n",
+    "**\\`Export a model with control flow to ONNX**\n",
+    "\n",
+    "Export a model with control flow to ONNX\n",
+    "========================================\n",
+    "\n",
+    "**Author**: [Xavier Dupré](https://github.com/xadupre)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Overview\n",
+    "========\n",
+    "\n",
+    "This tutorial demonstrates how to handle control flow logic while\n",
+    "exporting a PyTorch model to ONNX. It highlights the challenges of\n",
+    "exporting conditional statements directly and provides solutions to\n",
+    "circumvent them.\n",
+    "\n",
+    "Conditional logic cannot be exported into ONNX unless they refactored to\n",
+    "use `torch.cond`{.interpreted-text role=\"func\"}. Let\\'s start with a\n",
+    "simple model implementing a test.\n",
+    "\n",
+    "What you will learn:\n",
+    "\n",
+    "-   How to refactor the model to use `torch.cond`{.interpreted-text\n",
+    "    role=\"func\"} for exporting.\n",
+    "-   How to export a model with control flow logic to ONNX.\n",
+    "-   How to optimize the exported model using the ONNX optimizer.\n",
+    "\n",
+    "Prerequisites\n",
+    "-------------\n",
+    "\n",
+    "-   `torch >= 2.6`\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false
+   },
+   "outputs": [],
+   "source": [
+    "import torch"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Define the Models\n",
+    "=================\n",
+    "\n",
+    "Two models are defined:\n",
+    "\n",
+    "`ForwardWithControlFlowTest`: A model with a forward method containing\n",
+    "an if-else conditional.\n",
+    "\n",
+    "`ModelWithControlFlowTest`: A model that incorporates\n",
+    "`ForwardWithControlFlowTest` as part of a simple MLP. The models are\n",
+    "tested with a random input tensor to confirm they execute as expected.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false
+   },
+   "outputs": [],
+   "source": [
+    "class ForwardWithControlFlowTest(torch.nn.Module):\n",
+    "    def forward(self, x):\n",
+    "        if x.sum():\n",
+    "            return x * 2\n",
+    "        return -x\n",
+    "\n",
+    "\n",
+    "class ModelWithControlFlowTest(torch.nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "        self.mlp = torch.nn.Sequential(\n",
+    "            torch.nn.Linear(3, 2),\n",
+    "            torch.nn.Linear(2, 1),\n",
+    "            ForwardWithControlFlowTest(),\n",
+    "        )\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        out = self.mlp(x)\n",
+    "        return out\n",
+    "\n",
+    "\n",
+    "model = ModelWithControlFlowTest()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Exporting the Model: First Attempt\n",
+    "==================================\n",
+    "\n",
+    "Exporting this model using torch.export.export fails because the control\n",
+    "flow logic in the forward pass creates a graph break that the exporter\n",
+    "cannot handle. This behavior is expected, as conditional logic not\n",
+    "written using `torch.cond`{.interpreted-text role=\"func\"} is\n",
+    "unsupported.\n",
+    "\n",
+    "A try-except block is used to capture the expected failure during the\n",
+    "export process. If the export unexpectedly succeeds, an `AssertionError`\n",
+    "is raised.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false
+   },
+   "outputs": [],
+   "source": [
+    "x = torch.randn(3)\n",
+    "model(x)\n",
+    "\n",
+    "try:\n",
+    "    torch.export.export(model, (x,), strict=False)\n",
+    "    raise AssertionError(\"This export should failed unless PyTorch now supports this model.\")\n",
+    "except Exception as e:\n",
+    "    print(e)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Using `torch.onnx.export`{.interpreted-text role=\"func\"} with JIT\n",
+    "Tracing\n",
+    "\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\--\n",
+    "\n",
+    "When exporting the model using `torch.onnx.export`{.interpreted-text\n",
+    "role=\"func\"} with the dynamo=True argument, the exporter defaults to\n",
+    "using JIT tracing. This fallback allows the model to export, but the\n",
+    "resulting ONNX graph may not faithfully represent the original model\n",
+    "logic due to the limitations of tracing.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false
+   },
+   "outputs": [],
+   "source": [
+    "onnx_program = torch.onnx.export(model, (x,), dynamo=True) \n",
+    "print(onnx_program.model)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Suggested Patch: Refactoring with `torch.cond`{.interpreted-text\n",
+    "role=\"func\"}\n",
+    "\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\-\\--\n",
+    "\n",
+    "To make the control flow exportable, the tutorial demonstrates replacing\n",
+    "the forward method in `ForwardWithControlFlowTest` with a refactored\n",
+    "version that uses `torch.cond`{.interpreted-text role=\"func\"}\\`.\n",
+    "\n",
+    "Details of the Refactoring:\n",
+    "\n",
+    "Two helper functions (identity2 and neg) represent the branches of the\n",
+    "conditional logic: \\* `torch.cond`{.interpreted-text role=\"func\"}[ is\n",
+    "used to specify the condition and the two branches along with the input\n",
+    "arguments. \\* The updated forward method is then dynamically assigned to\n",
+    "the ]{.title-ref}[ForwardWithControlFlowTest]{.title-ref}\\` instance\n",
+    "within the model. A list of submodules is printed to confirm the\n",
+    "replacement.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false
+   },
+   "outputs": [],
+   "source": [
+    "def new_forward(x):\n",
+    "    def identity2(x):\n",
+    "        return x * 2\n",
+    "\n",
+    "    def neg(x):\n",
+    "        return -x\n",
+    "\n",
+    "    return torch.cond(x.sum() > 0, identity2, neg, (x,))\n",
+    "\n",
+    "\n",
+    "print(\"the list of submodules\")\n",
+    "for name, mod in model.named_modules():\n",
+    "    print(name, type(mod))\n",
+    "    if isinstance(mod, ForwardWithControlFlowTest):\n",
+    "        mod.forward = new_forward"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Let\\'s see what the FX graph looks like.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false
+   },
+   "outputs": [],
+   "source": [
+    "print(torch.export.export(model, (x,), strict=False))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Let\\'s export again.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false
+   },
+   "outputs": [],
+   "source": [
+    "onnx_program = torch.onnx.export(model, (x,), dynamo=True)  \n",
+    "print(onnx_program.model)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We can optimize the model and get rid of the model local functions\n",
+    "created to capture the control flow branches.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false
+   },
+   "outputs": [],
+   "source": [
+    "onnx_program.optimize()  \n",
+    "print(onnx_program.model)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Conclusion\n",
+    "==========\n",
+    "\n",
+    "This tutorial demonstrates the challenges of exporting models with\n",
+    "conditional logic to ONNX and presents a practical solution using\n",
+    "`torch.cond`{.interpreted-text role=\"func\"}. While the default exporters\n",
+    "may fail or produce imperfect graphs, refactoring the model\\'s logic\n",
+    "ensures compatibility and generates a faithful ONNX representation.\n",
+    "\n",
+    "By understanding these techniques, we can overcome common pitfalls when\n",
+    "working with control flow in PyTorch models and ensure smooth\n",
+    "integration with ONNX workflows.\n",
+    "\n",
+    "Further reading\n",
+    "===============\n",
+    "\n",
+    "The list below refers to tutorials that ranges from basic examples to\n",
+    "advanced scenarios, not necessarily in the order they are listed. Feel\n",
+    "free to jump directly to specific topics of your interest or sit tight\n",
+    "and have fun going through all of them to learn all there is about the\n",
+    "ONNX 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
+}