布局变换

参考：tvm/tests/python/contrib/test_onnx.py

import set_env

import numpy as np
import onnxruntime as rt

import onnx
import onnxruntime

import tvm
from tvm import relay
from tvm.contrib.target.onnx import to_onnx
from tvm.relay.testing import run_infer_type

def func_to_onnx(func, name):
    mod = tvm.IRModule()
    mod["main"] = func
    onnx_model = to_onnx(mod, {}, name, path=None)
    return onnx_model.SerializeToString()


def run_onnx(onnx_model, input_data):
    sess = rt.InferenceSession(onnx_model)
    input_names = {}
    for input, data in zip(sess.get_inputs(), input_data):
        input_names[input.name] = data
    output_names = [out.name for out in sess.get_outputs()]
    res = sess.run(output_names, input_names)
    return res


def run_relay(func, data_tuple, is_dyn=False):
    target = "llvm"
    dev = tvm.device("llvm", 0)
    kind = "graph" if not is_dyn else "vm"
    relay_res = relay.create_executor(kind, device=dev, target=target).evaluate(func)(*data_tuple)

    result = []
    relay_res = relay_res if isinstance(relay_res, list) else [relay_res]
    for res in relay_res:
        result.append(res.numpy())

    return result


def verify_results(relay_func, indata, test_name, rtol=1e-7, atol=0, is_dyn=False):
    relay_results = run_relay(relay_func, indata, is_dyn)
    onnx_results = run_onnx(func_to_onnx(relay_func, test_name), indata)

    for relay_res, onnx_res in zip(relay_results, onnx_results):
        np.testing.assert_allclose(relay_res, onnx_res, rtol=rtol, atol=atol)

def verify_layout_transform(dshape, src_layout, dst_layout, dtype="float32"):
    x = relay.var("x", relay.ty.TensorType(dshape, dtype))
    y = relay.layout_transform(x, src_layout, dst_layout)
    func = relay.Function([x], y)
    x_data = np.random.uniform(size=dshape).astype(dtype)
    verify_results(func, [x_data], "test_layout_transform", rtol=1e-5, atol=1e-5)

verify_layout_transform((1, 3, 8, 8), "NCHW", "NHWC")
verify_layout_transform((1, 8, 8, 3), "NHWC", "NCHW")