Compile the ONNX model

This article is translated from the English document Compile ONNX Models — tvm 0.13.dev0 documentation.
The author is Joshua Z. Zhang
. More TVM Chinese documents can be accessed → TVM Chinese Station.

This article will introduce how to deploy ONNX models with Relay.

First install the ONNX package. The most convenient way is to install the protobuf compiler:

pip install --user onnx onnxoptimizer

Or refer to the official website: https://github.com/onnx/onnx

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 the pre-trained ONNX model

The super-resolution model in the example below is exactly the same as the model in the ONNX tutorial , skip the PyTorch model building part, and download the saved ONNX model:

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")


# 现在磁盘上有 super_resolution.onnx 模型
onnx_model = onnx.load(model_path)

load test image

The model takes as input a single image of size 224x224 and outputs an image upscaled by a factor of 3 along each axis (ie size 672x672). To fit the input shape, the cat image is rescaled and converted to YCbCr. The super-resolution model is then applied to the luminance (Y) channel.

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

Usually ONNX models mix input values ​​with parameter values, and the input name is 1, check the model documentation to determine the complete input and parameter namespace.

Pass the shape dictionary to the relay.frontend.from_onnx method so that Relay knows which ONNX parameters are inputs and which are parameters, and provide a static definition of the input dimensions:

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()

Output result:

/workspace/python/tvm/relay/frontend/onnx.py:5785: UserWarning: Mismatched attribute type in ' : kernel_shape'

==> Context: Bad node spec for node. Name:  OpType: Conv
  warnings.warn(str(e))
/workspace/python/tvm/driver/build_module.py:268: UserWarning: target_host parameter is going to be deprecated. Please pass in tvm.target.Target(target, host=target_host) instead.
  "target_host parameter is going to be deprecated. "

Execute on TVM

dtype = "float32"
tvm_output = executor(tvm.nd.array(x.astype(dtype))).numpy()

View Results

Compare the input and output images side by side. The luminance channel Y is the output of the model. The chroma channels Cb and Cr are adjusted to match a simple bicubic algorithm, then the image is recombined and converted back to RGB.

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()

Output result:

/workspace/gallery/how_to/compile_models/from_onnx.py:120: DeprecationWarning: BICUBIC is deprecated and will be removed in Pillow 10 (2023-07-01). Use Resampling.BICUBIC instead.
  out_cb = img_cb.resize(out_y.size, Image.BICUBIC)
/workspace/gallery/how_to/compile_models/from_onnx.py:121: DeprecationWarning: BICUBIC is deprecated and will be removed in Pillow 10 (2023-07-01). Use Resampling.BICUBIC instead.
  out_cr = img_cr.resize(out_y.size, Image.BICUBIC)

Note that
the ONNX importer defines the model according to the dynamic shape by default when importing, and the compiler converts the model into a static shape when compiling. If it fails, there may still be dynamic operations in the model. Currently not all TVM kernels support dynamic shape. If you encounter dynamic kernel bugs, please submit an issue on discuss.tvm.apache.org.

This particular model was built with an older version of ONNX. During the import phase, the ONNX importer runs the ONNX validator (possibly throwing attribute type mismatch warnings). Since TVM supports many different ONNX versions, the Relay model is still valid.

Download the Python source code: from_onnx.py

Download Jupyter Notebook: from_onnx.ipynb