.. note:: :class: sphx-glr-download-link-note Click :ref:`here ` to download the full example code .. rst-class:: sphx-glr-example-title .. _sphx_glr_how_to_compile_models_from_onnx.py: Compile ONNX Models =================== **Author**: `Joshua Z. Zhang `_ This article is an introductory tutorial to deploy ONNX models with Relay. For us to begin with, ONNX package must be installed. A quick solution is to install protobuf compiler, and .. code-block:: bash pip install --user onnx onnxoptimizer or please refer to official site. https://github.com/onnx/onnx .. code-block:: default import onnx import numpy as np import tvm from tvm import te import tvm.relay as relay from tvm.contrib.download import download_testdata Load pretrained ONNX model --------------------------------------------- The example super resolution model used here is exactly the same model in onnx tutorial http://pytorch.org/tutorials/advanced/super_resolution_with_caffe2.html we skip the pytorch model construction part, and download the saved onnx model .. code-block:: default model_url = "".join( [ "https://gist.github.com/zhreshold/", "bcda4716699ac97ea44f791c24310193/raw/", "93672b029103648953c4e5ad3ac3aadf346a4cdc/", "super_resolution_0.2.onnx", ] ) model_path = download_testdata(model_url, "super_resolution.onnx", module="onnx") # now you have super_resolution.onnx on disk onnx_model = onnx.load(model_path) Load a test image --------------------------------------------- A single cat dominates the examples! This model takes a single input image of size 224x224 and outputs a scaled image that is 3x greater than the input along each axis, a 672x672 image. Re-scale the cat image to fit this input shape then convert to `YCbCr`. The super resolution model will then be applied to the luminance (`Y`) channel. .. code-block:: default from PIL import Image img_url = "https://github.com/dmlc/mxnet.js/blob/main/data/cat.png?raw=true" img_path = download_testdata(img_url, "cat.png", module="data") img = Image.open(img_path).resize((224, 224)) img_ycbcr = img.convert("YCbCr") # convert to YCbCr img_y, img_cb, img_cr = img_ycbcr.split() x = np.array(img_y)[np.newaxis, np.newaxis, :, :] Compile the model with relay --------------------------------------------- Typically ONNX models mix model input values with parameter values, with the input having the name `1`. This model dependent, and you should check with the documentation for your model to determine the full input and parameter name space. Passing in the shape dictionary to the `relay.frontend.from_onnx` method tells relay which ONNX parameters are inputs, and which are parameters, and provides a static definition of the input size. .. code-block:: default target = "llvm" input_name = "1" shape_dict = {input_name: x.shape} mod, params = relay.frontend.from_onnx(onnx_model, shape_dict) with tvm.transform.PassContext(opt_level=1): executor = relay.build_module.create_executor( "graph", mod, tvm.cpu(0), target, params ).evaluate() .. rst-class:: sphx-glr-script-out Out: .. code-block:: none /workspace/python/tvm/relay/frontend/onnx.py:4739: UserWarning: Mismatched attribute type in ' : kernel_shape' ==> Context: Bad node spec: input: "1" input: "2" output: "11" op_type: "Conv" attribute { name: "kernel_shape" ints: 5 ints: 5 } attribute { name: "strides" ints: 1 ints: 1 } attribute { name: "pads" ints: 2 ints: 2 ints: 2 ints: 2 } attribute { name: "dilations" ints: 1 ints: 1 } attribute { name: "group" i: 1 } warnings.warn(str(e)) Execute on TVM --------------------------------------------- .. code-block:: default dtype = "float32" tvm_output = executor(tvm.nd.array(x.astype(dtype))).numpy() Display results --------------------------------------------- We put input and output image neck to neck. The luminance channel, `Y` is the output from the model. The chroma channels `Cb` and `Cr` are resized to match with a simple bicubic algorithm. The image is then recombined and converted back to `RGB`. .. code-block:: default from matplotlib import pyplot as plt out_y = Image.fromarray(np.uint8((tvm_output[0, 0]).clip(0, 255)), mode="L") out_cb = img_cb.resize(out_y.size, Image.BICUBIC) out_cr = img_cr.resize(out_y.size, Image.BICUBIC) result = Image.merge("YCbCr", [out_y, out_cb, out_cr]).convert("RGB") canvas = np.full((672, 672 * 2, 3), 255) canvas[0:224, 0:224, :] = np.asarray(img) canvas[:, 672:, :] = np.asarray(result) plt.imshow(canvas.astype(np.uint8)) plt.show() .. image:: /how_to/compile_models/images/sphx_glr_from_onnx_001.png :class: sphx-glr-single-img Notes --------------------------------------------- By default, ONNX defines models in terms of dynamic shapes. The ONNX importer retains that dynamism upon import, and the compiler attempts to convert the model into a static shapes at compile time. If this fails, there may still be dynamic operations in the model. Not all TVM kernels currently support dynamic shapes, please file an issue on discuss.tvm.apache.org if you hit an error with dynamic kernels. This particular model was build using an older version of ONNX. During the import phase ONNX importer will run the ONNX verifier, which may throw a `Mismatched attribute type` warning. Because TVM supports a number of different ONNX versions, the Relay model will still be valid. .. _sphx_glr_download_how_to_compile_models_from_onnx.py: .. only :: html .. container:: sphx-glr-footer :class: sphx-glr-footer-example .. container:: sphx-glr-download :download:`Download Python source code: from_onnx.py ` .. container:: sphx-glr-download :download:`Download Jupyter notebook: from_onnx.ipynb ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_