Skip to content

Commit

Permalink
[Contrib][ONNX] Handle removal of onnx.utils.polish_model
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 committed Oct 1, 2021
1 parent de04d83 commit 18cbaec
Show file tree
Hide file tree
Showing 4 changed files with 33 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
25 changes: 23 additions & 2 deletions python/tvm/contrib/target/onnx.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,28 @@
ONNX_OPSET_VERSONS_SUPPORTED = [11]


def run_onnx_optimizer(onnx_model):
# ONNX Optimizer was removed from onnx in version 1.9 (see
# discussion in https://github.com/onnx/onnx/pull/2834).
try:
onnx_polish_model = onnx.utils.polish_model
except AttributeError:
pass
else:
return onnx_polish_model(onnx_model)

# If onnxoptimizer is installed, use it as a fallback case.
try:
import onnxoptimizer
except ImportError:
pass
else:
return onnxoptimizer.optimize(onnx_model)

# If neither are available, return the model as-is.
return model


def tvm_array_to_list(arr):
return tuple(x.value for x in arr)

Expand Down Expand Up @@ -881,8 +903,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 18cbaec

Please sign in to comment.