From 3c5a01f5870d87d73e1cbdc200a7d07e2560613f Mon Sep 17 00:00:00 2001
From: Lunderberg <Lunderberg@users.noreply.github.com>
Date: Sat, 2 Oct 2021 05:52:23 -0500
Subject: [PATCH] [Contrib][ONNX] Handle removal of onnx.utils.polish_model
 (#9178)

Onnx 1.9 removed optimizers from the core repository (see discussion
in https://github.com/onnx/onnx/pull/2834), including
onnx.utils.polish_model, breaking RelayToONNXConverter.  This PR add
onnxoptimizer.optimize as a fallback method if onnx.utils.polish_model
is unavailable.

Also updates tutorials/documentation to recommend installing
onnxoptimizer when installing onnx, because the current PyPI version
of onnx is above version 1.9.
---
 docker/install/ubuntu_install_onnx.sh         |  4 +++
 python/tvm/contrib/target/onnx.py             | 31 +++++++++++++++++--
 tutorials/frontend/from_onnx.py               |  2 +-
 .../get_started/tvmc_command_line_driver.py   |  9 +++---
 4 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/docker/install/ubuntu_install_onnx.sh b/docker/install/ubuntu_install_onnx.sh
index ef0bf1b012c6..37a9a1304927 100755
--- a/docker/install/ubuntu_install_onnx.sh
+++ b/docker/install/ubuntu_install_onnx.sh
@@ -22,6 +22,10 @@ set -o pipefail
 
 # We need to fix the onnx version because changing versions tends to break tests
 # TODO(mbrookhart): periodically update
+
+# onnx 1.9 removed onnx optimizer from the main repo (see
+# https://github.com/onnx/onnx/pull/2834).  When updating the CI image
+# to onnx>=1.9, onnxoptimizer should also be installed.
 pip3 install \
     onnx==1.8.1 \
     onnxruntime==1.7.0
diff --git a/python/tvm/contrib/target/onnx.py b/python/tvm/contrib/target/onnx.py
index 6f8aab23cde1..c26255fc5517 100644
--- a/python/tvm/contrib/target/onnx.py
+++ b/python/tvm/contrib/target/onnx.py
@@ -34,6 +34,34 @@
 ONNX_OPSET_VERSONS_SUPPORTED = [11]
 
 
+def run_onnx_optimizer(onnx_model):
+    """Run ONNX's optimization routines.
+
+    ONNX Optimizer was moved to an external library in
+    version 1.9.  Attempt to use the optimizer in onnx if
+    it is available, fall back to the standalone
+    onnxoptimizer otherwise, and return the model
+    unoptimized if neither are available.
+
+    """
+    try:
+        onnx_polish_model = onnx.utils.polish_model
+    except AttributeError:
+        pass
+    else:
+        return onnx_polish_model(onnx_model)
+
+    try:
+        # pylint: disable=import-outside-toplevel
+        import onnxoptimizer
+    except ImportError:
+        pass
+    else:
+        return onnxoptimizer.optimize(onnx_model)
+
+    return model
+
+
 def tvm_array_to_list(arr):
     return tuple(x.value for x in arr)
 
@@ -881,8 +909,7 @@ def convert_to_onnx(self, func):
         self.visit(func)
         self._add_output(self._node_dict[self.last_node])
         model = self._mc.make_model()
-        polished_model = onnx.utils.polish_model(model)
-        return polished_model
+        return run_onnx_optimizer(model)
 
     def visit(self, expr):
         self._node_count += 1
diff --git a/tutorials/frontend/from_onnx.py b/tutorials/frontend/from_onnx.py
index fd51d7a76992..586c811aa627 100644
--- a/tutorials/frontend/from_onnx.py
+++ b/tutorials/frontend/from_onnx.py
@@ -27,7 +27,7 @@
 
 .. code-block:: bash
 
-    pip install onnx --user
+    pip install --user onnx onnxoptimizer
 
 or please refer to official site.
 https://github.com/onnx/onnx
diff --git a/tutorials/get_started/tvmc_command_line_driver.py b/tutorials/get_started/tvmc_command_line_driver.py
index c729b86a3245..ea3254054ecf 100644
--- a/tutorials/get_started/tvmc_command_line_driver.py
+++ b/tutorials/get_started/tvmc_command_line_driver.py
@@ -104,10 +104,11 @@
 ################################################################################
 # .. note:: Adding ONNX Support to TVM
 #
-#    TVM relies on the ONNX python library being available on your system. You
-#    can install ONNX using the command ``pip3 install --user onnx``. You may
-#    remove the ``--user`` option if you have root access and want to install
-#    ONNX globally.
+#    TVM relies on the ONNX python library being available on your system. You can
+#    install ONNX using the command ``pip3 install --user onnx onnxoptimizer``. You
+#    may remove the ``--user`` option if you have root access and want to install
+#    ONNX globally.  The ``onnxoptimizer`` dependency is optional, and is only used
+#    for ``onnx>=1.9``.
 #
 
 ################################################################################