ONNX預訓練模型加載

tvm官網中，對從ONNX預訓練模型中加載模型的教程說明

教程來自於：https://docs.tvm.ai/tutorials/frontend/from_onnx.html#sphx-glr-tutorials-frontend-from-onnx-py

首先我對教程進行了一些修改，很多東西沒有必要，比如不是每次都需要從網上下載圖片和模型，super_resolution.onnx和cat.png都預先下載到了文件同目錄下，

同時，最新版本的tvm中不支持Python2.7，我沒有編譯llvm，所以我把我的設置都改到了cuda上，在24行和32行有體現,注意最新版本

import onnx
import numpy as np
import tvm
import tvm.relay as relay
# from tvm.contrib.download import download_testdata

# 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('super_resolution.onnx')

from PIL import Image
# img_url = 'https://github.com/dmlc/mxnet.js/blob/master/data/cat.png?raw=true'
# img_path = download_testdata(img_url, 'cat.png', module='data')
img_path = 'cat.png'
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, :, :]

target = 'cuda'

input_name = '1'
shape_dict = {input_name: x.shape}
sym, params = relay.frontend.from_onnx(onnx_model, shape_dict)
print(sym)

with relay.build_config(opt_level=1):
    intrp = relay.build_module.create_executor('graph', sym, tvm.gpu(0), target)

dtype = 'float32'
tvm_output = intrp.evaluate(sym)(tvm.nd.array(x.astype(dtype)), **params).asnumpy()

第28行有一個從模型加載的函數from_onnx

官方的解釋：tvm.relay.frontend.from_onnx(model, shape=None, dtype='float32')

Convert a ONNX model into an equivalent Relay Function.

ONNX graphs are represented as Python Protobuf objects. The companion parameters will be handled automatically. However, the input names from onnx graph is vague, mixing inputs and network weights/bias such as “1”, “2”… For convenience, we rename the real input names to “input_0”, “input_1”… And renaming parameters to “param_0”, “param_1”…

Parameters:	model (protobuf object) – ONNX ModelProto after ONNX v1.1.0 shape (dict of str to tuple, optional) – The input shape to the graph dtype (str or dict of str to str) – The input types to the graph
Returns:	sym (tvm.relay.expr.Function) – Compatible relay function params (dict of str to tvm.NDArray) – The parameter dict to be used by relay

看返回值，sym是relay Function，在后邊加一個print(sym)輸出，可以看到圖這一級的IR

fn (%v1: Tensor[(1, 1, 224, 224), float32], %v2: Tensor[(64, 1, 5, 5), float32], %v3: Tensor[(64,), float32], %v4: Tensor[(64, 64, 3, 3), float32], %v5: Tensor[(64,), float32], %v6: Tensor[(32, 64, 3, 3), float32], %v7: Tensor[(32,), float32], %v8: Tensor[(9, 32, 3, 3), float32], %v9: Tensor[(9,), float32]) {
  %0 = nn.conv2d(%v1, %v2, padding=[2, 2], kernel_size=[5, 5])
  %1 = expand_dims(%v3, axis=1, num_newaxis=2)
  %2 = add(%0, %1)
  %3 = nn.relu(%2)
  %4 = nn.conv2d(%3, %v4, padding=[1, 1], kernel_size=[3, 3])
  %5 = expand_dims(%v5, axis=1, num_newaxis=2)
  %6 = add(%4, %5)
  %7 = nn.relu(%6)
  %8 = nn.conv2d(%7, %v6, padding=[1, 1], kernel_size=[3, 3])
  %9 = expand_dims(%v7, axis=1, num_newaxis=2)
  %10 = add(%8, %9)
  %11 = nn.relu(%10)
  %12 = nn.conv2d(%11, %v8, padding=[1, 1], kernel_size=[3, 3])
  %13 = expand_dims(%v9, axis=1, num_newaxis=2)
  %14 = add(%12, %13)
  %15 = reshape(%14, newshape=[1, 1, 3, 3, 224, 224])
  %16 = transpose(%15, axes=[0, 1, 4, 2, 5, 3])
  reshape(%16, newshape=[1, 1, 672, 672])
}