add torch.cond

Author: titaiwangms

beginner_source/onnx/README.txt

@@ -12,3 +12,7 @@ ONNX
 3. onnx_registry_tutorial.py
     Extending the ONNX Registry
     https://pytorch.org/tutorials/beginner/onnx/onnx_registry_tutorial.html
+
+4. export_control_flow_model_to_onnx_tutorial.py
+    Export a Pytorch model with a test to ONNX
+    https://pytorch.org/tutorials/beginner/onnx/export_control_flow_model_to_onnx_tutorial.html

beginner_source/onnx/export_control_flow_model_to_onnx_tutorial.py

@@ -0,0 +1,104 @@
+# -*- coding: utf-8 -*-
+"""
+**Introduction to ONNX** ||
+`Exporting a PyTorch model to ONNX <export_simple_model_to_onnx_tutorial.html>`_ ||
+`Extending the ONNX Registry <onnx_registry_tutorial.html>`_
+
+Export a control flow model to ONNX
+==========================================
+
+**Author**: `Xavier Dupré <https://github.com/xadupre>`_.
+
+Tests cannot be exported into ONNX unless they refactored
+to use :func:`torch.cond`. Let's start with a simple model
+implementing a test.
+"""
+
+from onnx.printer import to_text
+import torch
+
+class ForwardWithControlFlowTest(torch.nn.Module):
+    def forward(self, x):
+        if x.sum():
+            return x * 2
+        return -x
+
+
+class ModelWithControlFlowTest(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.mlp = torch.nn.Sequential(
+            torch.nn.Linear(3, 2),
+            torch.nn.Linear(2, 1),
+            ForwardWithControlFlowTest(),
+        )
+
+    def forward(self, x):
+        out = self.mlp(x)
+        return out
+
+
+model = ModelWithControlFlowTest()
+
+# %%
+# Let's check it runs.
+x = torch.randn(3)
+model(x)
+
+# %%
+# As expected, it does not export.
+try:
+    torch.export.export(model, (x,))
+    raise AssertionError("This export should failed unless pytorch now supports this model.")
+except Exception as e:
+    print(e)
+
+# %%
+# It does export with :func:`torch.onnx.export` because
+# it uses JIT to trace the execution.
+# But the model is not exactly the same as the initial model.
+ep = torch.onnx.export(model, (x,), dynamo=True)
+print(to_text(ep.model_proto))
+
+
+# %%
+# Suggested Patch
+# +++++++++++++++
+#
+# Let's avoid the graph break by replacing the forward.
+
+
+def new_forward(x):
+    def identity2(x):
+        return x * 2
+
+    def neg(x):
+        return -x
+
+    return torch.cond(x.sum() > 0, identity2, neg, (x,))
+
+
+print("the list of submodules")
+for name, mod in model.named_modules():
+    print(name, type(mod))
+    if isinstance(mod, ForwardWithControlFlowTest):
+        mod.forward = new_forward
+
+# %%
+# Let's see what the fx graph looks like.
+
+print(torch.export.export(model, (x,)).graph)
+
+# %%
+# Let's export again.
+
+ep = torch.onnx.export(model, (x,), dynamo=True)
+print(to_text(ep.model_proto))
+
+
+# %%
+# Let's optimize to see a small model.
+
+ep = torch.onnx.export(model, (x,), dynamo=True)
+ep.optimize()
+print(to_text(ep.model_proto))

beginner_source/onnx/onnx_toc.txt

@@ -1,2 +1,3 @@
 | 1. `Exporting a PyTorch model to ONNX <export_simple_model_to_onnx_tutorial.html>`_
 | 2. `Extending the ONNX registry <onnx_registry_tutorial.html>`_
+| 3. `Export a Pytorch model with a test to ONNX <export_control_flow_model_to_onnx_tutorial>`_