Skip to content

Commit

Permalink
[Contrib][ONNX] Handle removal of onnx.utils.polish_model (apache#9178)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
Lunderberg authored and ylc committed Jan 13, 2022
1 parent b932916 commit 3c5a01f
Show file tree
Hide file tree
Showing 4 changed files with 39 additions and 7 deletions.
4 changes: 4 additions & 0 deletions docker/install/ubuntu_install_onnx.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
31 changes: 29 additions & 2 deletions python/tvm/contrib/target/onnx.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion tutorials/frontend/from_onnx.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
9 changes: 5 additions & 4 deletions tutorials/get_started/tvmc_command_line_driver.py
Original file line number Diff line number Diff line change
Expand Up @@ -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``.
#

################################################################################
Expand Down

0 comments on commit 3c5a01f

Please sign in to comment.