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