Compilar el modelo ONNX

Este artículo está traducido del documento en inglés Compile ONNX Models — tvm 0.13.dev0 documentation.
El autor es Joshua Z. Zhang
. Se puede acceder a más documentos en chino de TVM → Estación china de TVM.

Este artículo presentará cómo implementar modelos ONNX con Relay.

Primero instale el paquete ONNX, la forma más conveniente es instalar el compilador protobuf:

pip install --user onnx onnxoptimizer

O consulte el sitio web oficial: 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

Cargue el modelo ONNX preentrenado

El modelo de superresolución en el siguiente ejemplo es exactamente el mismo que el modelo en el tutorial de ONNX , omita la parte de construcción del modelo PyTorch y descargue el modelo ONNX guardado:

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)

imagen de prueba de carga

El modelo toma como entrada una sola imagen de tamaño 224x224 y genera una imagen ampliada por un factor de 3 a lo largo de cada eje (es decir, tamaño 672x672). Para ajustarse a la forma de entrada, la imagen del gato se vuelve a escalar y se convierte a YCbCr. A continuación, se aplica el modelo de superresolución al canal de luminancia (Y).

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 el modelo con Relay

Por lo general, los modelos ONNX combinan valores de entrada con valores de parámetros, y el nombre de entrada es 1; consulte la documentación del modelo para determinar el espacio de nombres de entrada y parámetro completo.

Pase el diccionario de formas al método relay.frontend.from_onnx para que Relay sepa qué parámetros ONNX son entradas y cuáles son parámetros, y proporcione una definición estática de las dimensiones de entrada:

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

Resultado de salida:

/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. "

Ejecutar en TVM

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

Ver resultados

Compare las imágenes de entrada y salida una al lado de la otra. El canal de luminancia Y es la salida del modelo. Los canales de croma Cb y Cr se ajustan para que coincidan con un algoritmo bicúbico simple, luego la imagen se recombina y se vuelve a convertir a 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()

Resultado de salida:

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

Tenga en cuenta que
el importador ONNX define el modelo de acuerdo con la forma dinámica de forma predeterminada al importar, y el compilador convierte el modelo en una forma estática al compilar. Si falla, todavía puede haber operaciones dinámicas en el modelo. Actualmente, no todos los kernels de TVM son compatibles con la forma dinámica. Si encuentra errores de kernel dinámicos, envíe un problema en debate.tvm.apache.org.

Este modelo en particular fue construido con una versión anterior de ONNX. Durante la fase de importación, el importador ONNX ejecuta el validador ONNX (posiblemente emitiendo advertencias de discrepancia de tipo de atributo). Dado que TVM admite muchas versiones diferentes de ONNX, el modelo de retransmisión sigue siendo válido.

Descarga el código fuente de Python: from_onnx.py

Descargar Jupyter Notebook: from_onnx.ipynb