[Contrib][ONNX] Handle removal of onnx.utils.polish_model (#9178)

Onnx 1.9 removed optimizers from the core repository (see discussion
in onnx/onnx#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

@@ -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

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

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

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``.
 #
 
 ################################################################################