def from_auto_ml(model, shape=None, dtype="float32", func_name="transform"):
"""
Import scikit-learn model to Relay.
"""
try:
import sklearn # pylint: disable=unused-import
except ImportError as e:
raise ImportError(
"Unable to import scikit-learn which is required {}".format(e))
if func_name == "transform":
inexpr_float = _expr.var("input_float", shape=shape, dtype=dtype)
inexpr_string = _expr.var("input_string", shape=shape, dtype=dtype)
inexpr = [inexpr_float, inexpr_string]
if type(model.feature_transformer.steps[0]
[1]).__name__ != "ColumnTransformer":
raise NameError(
"The First Transformer must be an ColumnTransformer, but {} is given"
.format(type(model.feature_transformer.steps[0][1]).__name__))
outexpr = inexpr
for _, transformer in model.feature_transformer.steps:
outexpr = sklearn_op_to_relay(transformer, outexpr, shape, dtype,
func_name, None)
else:
inexpr = _expr.var("input", shape=shape, dtype=dtype)
transformer = model.target_transformer
outexpr = sklearn_op_to_relay(transformer, inexpr, shape, dtype,
func_name, None)
func = _function.Function(analysis.free_vars(outexpr), outexpr)
return IRModule.from_expr(func), []
def from_translated_layer(self, layer, shape_dict):
"""Construct the TVM relay expression from PaddlePaddle TranslatedLayer."""
self.shape_dict = shape_dict
program = layer.program()
parameters = dict()
for param in layer.parameters():
parameters[param.name] = np.array(param.value().get_tensor())
self.check_unsupported_ops(program)
self.extract_parameters(program, parameters)
input_specs = layer._input_spec()
self.ops_to_relay(program, input_specs)
output_names = [x.name for x in layer._output_spec()]
outputs = [self.nodes[name] for name in output_names]
outputs = outputs[0] if len(outputs) == 1 else _expr.Tuple(outputs)
free_vars = analysis.free_vars(outputs)
func = _function.Function(free_vars, outputs)
mod = IRModule.from_expr(func)
# remove unused parameters
final_params = dict()
for var in free_vars:
if var.name_hint in self.params:
final_params[var.name_hint] = self.params[var.name_hint]
self.params = final_params
return mod, self.params
def from_darknet(self):
"""To convert the darknet symbol to relay functions."""
for i in range(self._net.n):
layer = self._net.layers[i]
need_skip, sym = self._preproc_layer(layer, i)
if need_skip:
continue
processed, sym = self._handle_darknet_rnn_layers(i, sym)
if processed:
continue
attr = self._get_darknet_attrs(layer, i)
op_name = self._get_opname(layer)
prefix = _get_params_prefix(op_name, i)
params = self._get_darknet_params(self._net.layers[i], prefix)
sym = _darknet_convert_symbol(op_name, _as_list(sym), params, attr,
prefix)
if params:
self._tvmparams.update(params)
self._sym_array[i] = sym
self._make_outlist(sym, prefix, layer, i)
outputs = _as_list(sym) + self._outs
outputs = outputs[0] if len(outputs) == 1 else _expr.Tuple(outputs)
sym = _function.Function(analysis.free_vars(outputs), outputs)
return IRModule.from_expr(sym), self._tvmparams
def create_executor(kind="debug", mod=None, device=None, target="llvm", params=None):
"""Factory function to create an executor.
Example
-------
.. code-block:: python
import tvm.relay
import numpy as np
x = tvm.relay.var("x", tvm.relay.TensorType([1], dtype="float32"))
expr = tvm.relay.add(x, tvm.relay.Constant(tvm.nd.array(np.array([1], dtype="float32"))))
tvm.relay.create_executor(
kind="vm", mod=tvm.IRModule.from_expr(tvm.relay.Function([x], expr))
).evaluate()(np.array([2], dtype="float32"))
# returns `array([3.], dtype=float32)`
Parameters
----------
kind : str
The type of executor. Avaliable options are `debug` for the
interpreter, `graph` for the graph executor, and `vm` for the virtual
machine.
mod : :py:class:`~tvm.IRModule`
The Relay module containing collection of functions
device : :py:class:`Device`
The device to execute the code.
target : :py:class:`tvm.Target`
The corresponding context
params : dict of str to NDArray
Input parameters to the graph that do not change
during inference time.
Returns
-------
executor : :py:class:`~tvm.relay.backend.interpreter.Executor`
"""
if mod is None:
mod = IRModule()
if device is not None:
assert device.device_type == _nd.device(str(target), 0).device_type
else:
device = _nd.device(str(target), 0)
if params is not None:
mod = IRModule.from_expr(bind_params_by_name(mod["main"], params))
if isinstance(target, str):
target = Target(target)
if kind == "debug":
return _interpreter.Interpreter(mod, device, target)
if kind == "graph":
return GraphExecutor(mod, device, target)
if kind == "vm":
return VMExecutor(mod, device, target)
raise RuntimeError("unknown execution strategy: {0}".format(kind))
def _interp_wrapper(*args, **kwargs):
if expr is None:
args = self._convert_args(self.mod["main"], args, kwargs)
else:
args = self._convert_args(expr, args, kwargs)
relay_args = []
for arg in args:
relay_args.append(_arg_to_ast(self.mod, arg))
# Set the entry function for the module.
if expr is None:
pass
elif isinstance(expr, GlobalVar):
self.mod["main"] = self.mod[expr]
else:
assert isinstance(expr, Function)
func = Function([], Call(expr, relay_args))
relay_args = []
if self.mod:
self.mod["main"] = func
else:
self.mod = IRModule.from_expr(func)
mod = self.optimize()
opt_expr = Call(mod["main"], relay_args)
return _intrp(opt_expr)
def infer_value(input_val, params, mod=None):
"""A hack for getting the value of an expression by evaluating a
portion of the relay graph. This is often needed for functions that
whose output shape depends on the value of a tensor.
"""
# Check that all free variables have associated parameters.
assert all(
var.name_hint in params.keys() for var in analysis.free_vars(input_val)
), "All inputs to infer must be available in params."
try:
# TODO(kevinthesun): Use VM for all cases.
# pylint: disable=import-outside-toplevel
from tvm.contrib import graph_runtime
func = _function.Function(analysis.free_vars(input_val), input_val)
with tvm.transform.PassContext(opt_level=0):
lib = tvm.relay.build(func, target="llvm", params=params)
ctx = tvm.cpu(0)
m = graph_runtime.GraphModule(lib["default"](ctx))
m.run()
return m.get_output(0)
except Exception:
if isinstance(mod, IRModule):
mod["main"] = _function.Function(analysis.free_vars(input_val), input_val)
else:
mod = IRModule.from_expr(input_val)
exc = tvm.relay.create_executor("debug", mod=mod, ctx=tvm.cpu(), target="llvm")
inputs = []
for param in mod["main"].params:
inputs.append(params[param.name_hint])
result = exc.evaluate()(*inputs)
return result
def _set_params(mod, input_scale_func, weight_scale_func):
quantize_op = _op.get("relay.op.annotation.simulated_quantize")
cfg = quantize.current_qconfig()
const_params = {}
def visit_func(expr):
'''visitor function for traverse'''
if isinstance(expr, _expr.Call) and expr.op == quantize_op:
_, ndom_scale, nclip_min, nclip_max = expr.args
attrs = expr.attrs
kind = attrs.kind
nbit = cfg.get_nbit_by_kind(kind)
valid_bit = nbit - attrs.sign
# set scale
if kind == quantize.QAnnotateKind.WEIGHT:
assert isinstance(expr.args[0], _expr.Constant)
scale = weight_scale_func(expr)
else:
scale = input_scale_func(expr)
def _make_const(val):
return _expr.const(val, 'float32')
valid_range = 2**valid_bit
const_params[ndom_scale] = _make_const(scale / valid_range)
const_params[nclip_min] = _make_const(- (valid_range - 1))
const_params[nclip_max] = _make_const((valid_range - 1))
func = mod['main']
_analysis.post_order_visit(func, visit_func)
func = _expr.bind(func, const_params)
return IRModule.from_expr(func)
def from_program(self, program, shape_dict, scope):
"""Construct the TVM relay expression from PaddlePaddle program."""
self.shape_dict = shape_dict
if scope is None:
import paddle
scope = paddle.fluid.global_scope()
self.check_unsupported_ops(program)
self.extract_parameters(program, scope)
self.ops_to_relay(program)
output_names = list()
for block in program.blocks:
for op in block.ops:
if op.type == "fetch":
output_names.append(op.input("X")[0])
outputs = [self.nodes[name] for name in output_names]
outputs = outputs[0] if len(outputs) == 1 else _expr.Tuple(outputs)
free_vars = analysis.free_vars(outputs)
func = _function.Function(free_vars, outputs)
mod = IRModule.from_expr(func)
return mod, self.params
def check(dim, axis, nstep):
eps = 0.01
ttype1 = rly.TensorType(tuple(10 for i in range(dim)), dtype)
ttype2 = rly.TensorType((10, ), dtype)
x = rly.var("x", ttype1)
beta = rly.var("beta", ttype2)
gamma = rly.var("gamma", ttype2)
moving_var = rly.var("moving_var", ttype2)
moving_mean = rly.var("moving_mean", ttype2)
y1, y2 = x, x
for _ in range(nstep):
y1, _, _ = rly.nn.batch_norm(y1 + rly.const(1, dtype),
gamma,
beta,
moving_mean,
moving_var,
epsilon=eps,
axis=axis)
y1 = rly.nn.dropout(y1)
y2 = simple_bn(y2 + rly.const(1, dtype),
gamma,
beta,
moving_mean,
moving_var,
epsilon=eps,
axis=axis,
shape=ttype1.shape)
mod = IRModule.from_expr(y1)
simplify = SimplifyInference()
mod = simplify(mod)
y1 = mod["main"].body
assert rly.analysis.graph_equal(y1, y2)
def infer_type(node, mod=None):
"""A method to infer the type of an intermediate node in the relay graph."""
new_mod = IRModule.from_expr(node)
if mod is not None:
new_mod.update(mod)
new_mod = _transform.InferType()(new_mod)
entry = new_mod["main"]
return entry if isinstance(node, _function.Function) else entry.body
def optimize(mod, target=None, params=None):
"""Helper function that optimizes a Relay module.
Parameters
----------
mod : :py:class:`~tvm.IRModule`
The module to build. Using relay.Function is deprecated.
target : None, or any multi-target like object, see Target.canon_multi_target
For homogeneous compilation, the unique build target.
For heterogeneous compilation, a dictionary or list of possible build targets.
Defaults to the current target in the environment if None.
params : dict of str to NDArray
Input parameters to the graph that do not change
during inference time. Used for constant folding.
Returns
-------
mod : :py:class:`~tvm.IRModule`
The optimized relay module.
params : dict
The parameters of the final graph.
"""
if not isinstance(mod, (IRModule, _function.Function)):
raise ValueError("Type of input parameter mod must be tvm.IRModule")
if isinstance(mod, _function.Function):
if params:
mod = bind_params_by_name(mod, params)
mod = IRModule.from_expr(mod)
warnings.warn(
"Please use input parameter mod (tvm.IRModule) "
"instead of deprecated parameter func (tvm.relay.function.Function)",
DeprecationWarning,
)
raw_targets = Target.canon_multi_target_and_host(
Target.target_or_current(target))
# If current dispatch context is fallback context (the default root context),
# then load pre-tuned parameters from TopHub
if isinstance(autotvm.DispatchContext.current, autotvm.FallbackContext):
tophub_context = autotvm.tophub.context(raw_targets)
else:
tophub_context = autotvm.utils.EmptyContext()
with tophub_context:
bld_mod = BuildModule()
mod, params = bld_mod.optimize(mod, target=raw_targets, params=params)
return mod, params
def optimize(mod, target=None, params=None):
"""Helper function that optimizes a Relay module.
Parameters
----------
mod : :py:class:`~tvm.IRModule`
The module to build. Using relay.Function is deprecated.
target : str, :any:`tvm.target.Target`, or dict of str(i.e. device/context
name) to str/tvm.target.Target, optional
For heterogeneous compilation, it is a dictionary indicating context to
target mapping. For homogeneous compilation, it is a build target.
params : dict of str to NDArray
Input parameters to the graph that do not change
during inference time. Used for constant folding.
Returns
-------
mod : :py:class:`~tvm.IRModule`
The optimized relay module.
params : dict
The parameters of the final graph.
"""
if not isinstance(mod, (IRModule, _function.Function)):
raise ValueError("Type of input parameter mod must be tvm.IRModule")
if isinstance(mod, _function.Function):
if params:
mod = bind_params_by_name(mod, params)
mod = IRModule.from_expr(mod)
warnings.warn(
"Please use input parameter mod (tvm.IRModule) "
"instead of deprecated parameter func (tvm.relay.function.Function)",
DeprecationWarning,
)
target = _update_target(target)
# If current dispatch context is fallback context (the default root context),
# then load pre-tuned parameters from TopHub
if isinstance(autotvm.DispatchContext.current, autotvm.FallbackContext):
tophub_context = autotvm.tophub.context(list(target.values()))
else:
tophub_context = autotvm.utils.EmptyContext()
with tophub_context:
bld_mod = BuildModule()
mod, params = bld_mod.optimize(mod, target, params)
return mod, params
def visit_call(self, call):
# Keep track of our current depth and layer count
# so we can know whether to skip this layer or not.
current_depth = self.depth_count
current_layer = self.valid_op_count - current_depth - 1
if call.op in self.valid_ops:
self.depth_count += 1
# Visit current call operation
new_fn = self.visit(call.op)
# Visit current arguments
args = []
for arg in call.args:
args.append(self.visit(arg))
self.depth_count = current_depth
# Downcast this op if its the correct type and not skipped.
if call.op in self.valid_ops and current_layer not in self.skip_layers:
# Recast inputs to specified type.
args = [self.visit(arg) for arg in call.args]
new_args = list()
for arg in args:
new_args.append(relay.cast(arg, dtype=self.dtype))
# If out_dtype is in the attributes, we need to update it.
orig_dtype = None
if call.attrs is not None and "out_dtype" in call.attrs.keys():
new_attr_dict = {}
for attr in call.attrs.keys():
attr_value = call.attrs[attr]
if isinstance(attr_value, tvm.ir.container.Array):
attr_value = tuple(attr_value)
new_attr_dict[str(attr)] = attr_value
new_attr_dict["out_dtype"] = self.out_dtype
attr_type = str(call.attrs).split("(")[0]
new_attrs = tvm.ir.make_node(attr_type, **new_attr_dict)
if call.attrs["out_dtype"] != "":
orig_dtype = call.attrs["out_dtype"]
else:
new_attrs = call.attrs
if orig_dtype is None:
# Perform type inference to determine the original type.
new_mod = IRModule.from_expr(call)
new_mod = InferType()(new_mod)
checked_arg = new_mod["main"].body
orig_dtype = checked_arg.checked_type.dtype
# Recast the output for compatibility with other graph operations.
return relay.cast(Call(new_fn, new_args, new_attrs), orig_dtype)
# Otherwise return the unchanged call.
return Call(new_fn, args, call.attrs)
def preprocess_mod(mod, params):
if not isinstance(mod, (IRModule, _function.Function)):
raise ValueError("Type of input parameter mod must be tvm.IRModule")
if isinstance(mod, _function.Function):
if params:
mod = bind_params_by_name(mod, params)
mod = IRModule.from_expr(mod)
warnings.warn(
"Please use input parameter mod (tvm.IRModule) "
"instead of deprecated parameter mod (tvm.relay.function.Function)",
DeprecationWarning,
)
return mod
def from_sklearn(model, shape=None, dtype="float32", func_name="transform", columns=None):
"""
Import scikit-learn model to Relay.
"""
try:
import sklearn # pylint: disable=unused-import
except ImportError as e:
raise ImportError("Unable to import scikit-learn which is required {}".format(e))
if type(model).__name__ == "ColumnTransformer":
raise NameError("ColumnTransformer is not supported for single op compilation.")
inexpr = _expr.var("input", shape=shape, dtype=dtype)
outexpr = sklearn_op_to_relay(model, inexpr, shape, dtype, func_name, columns)
func = _function.Function(analysis.free_vars(outexpr), outexpr)
return IRModule.from_expr(func), []
def infer_type(node, mod=None):
"""A method to infer the type of an intermediate node in the relay graph."""
if isinstance(mod, IRModule):
mod["main"] = _function.Function(tvm.relay.analysis.free_vars(node), node)
mod = _transform.InferType()(mod)
entry = mod["main"]
ret = entry.body
else:
new_mod = IRModule.from_expr(node)
if mod is not None:
new_mod.update(mod)
new_mod = _transform.InferType()(new_mod)
entry = new_mod["main"]
ret = entry if isinstance(node, _function.Function) else entry.body
return ret
def test_batchnorm():
Device.load("test_config.json")
if skip_runtime_test():
return
device = Device()
np.random.seed(0)
dtype = "float32"
in_shape = (1, 8, 64, 64)
channels = 8
input_arr = tvm.nd.array(np.random.uniform(-1, 1, in_shape).astype(dtype))
inp = relay.var("a", shape=in_shape, dtype=dtype)
gamma_arr = tvm.nd.array(np.random.uniform(-1, 1, (channels)).astype(dtype))
beta_arr = tvm.nd.array(np.random.uniform(-1, 1, (channels)).astype(dtype))
gamma = relay.const(gamma_arr, dtype)
beta = relay.const(beta_arr, dtype)
mean_arr = tvm.nd.array(np.mean(input_arr.asnumpy(), axis=(0, 2, 3), keepdims=False))
mean = relay.const(mean_arr)
variance_arr = tvm.nd.array(np.var(input_arr.asnumpy(), axis=(0, 2, 3), keepdims=False))
variance = relay.const(variance_arr)
params = {}
func = relay.nn.batch_norm(inp, gamma, beta, mean, variance, axis=1, epsilon=0.0001)[0]
mod = IRModule.from_expr(func)
inputs = {
"a": input_arr,
}
opencl_out = build_and_run(mod, inputs, 1, params, device, enable_clml=False)[0]
clml_out = build_and_run(mod, inputs, 1, params, device, enable_clml=True)[0]
tvm.testing.assert_allclose(
clml_out[0].asnumpy(), opencl_out[0].asnumpy(), rtol=1e-5, atol=1e-5
)
def _set_params(mod, input_scale_func, weight_scale_func):
quantize_op = _op.get("relay.op.annotation.simulated_quantize")
cfg = quantize.current_qconfig()
const_params = {}
def visit_func(expr):
"""visitor function for traverse"""
if isinstance(expr, _expr.Call) and expr.op == quantize_op:
_, ndom_scale, nclip_min, nclip_max = expr.args
attrs = expr.attrs
kind = attrs.kind
nbit = cfg.get_nbit_by_kind(kind)
valid_bit = nbit - attrs.sign
# set scale
if kind == quantize.QAnnotateKind.WEIGHT:
assert isinstance(expr.args[0], _expr.Constant)
scale = weight_scale_func(expr)
else:
scale = input_scale_func(expr)
def _make_const(val):
return _expr.const(val, "float32")
valid_range = 2**valid_bit
const_params[ndom_scale] = _make_const(scale / valid_range)
const_params[nclip_min] = _make_const(-(valid_range - 1))
const_params[nclip_max] = _make_const((valid_range - 1))
main_func = mod["main"]
_analysis.post_order_visit(main_func, visit_func)
main_func = _expr.bind(main_func, const_params)
func_dict = {}
for global_var, func in mod.functions.items():
if global_var.name_hint != "main":
func_dict[global_var] = func
return IRModule.from_expr(main_func, func_dict)
def fold_constant(node, mod=None):
if mod is None:
mod = IRModule.from_expr(node)
return _transform.FoldConstantExpr(node, mod)
def build(ir_mod,
target=None,
target_host=None,
params=None,
mod_name="default"):
# fmt: off
# pylint: disable=line-too-long
"""Helper function that builds a Relay function to run on TVM graph executor.
Parameters
----------
ir_mod : :py:class:`~tvm.IRModule`
The IR module to build. Using relay.Function is deprecated.
target : str, :any:`tvm.target.Target`, or dict of str(i.e. device/context name) to str/tvm.target.Target, optional
For heterogeneous compilation, it is a dictionary indicating context to
target mapping. For homogeneous compilation, it is a build target.
target_host : str or :any:`tvm.target.Target`, optional
Host compilation target, if target is device.
When TVM compiles device specific program such as CUDA,
we also need host(CPU) side code to interact with the driver
setup the dimensions and parameters correctly.
target_host is used to specify the host side codegen target.
By default, llvm is used if it is enabled,
otherwise a stackvm intepreter is used.
params : dict of str to NDArray
Input parameters to the graph that do not change
during inference time. Used for constant folding.
mod_name: Optional[str]
The module name we will build
Returns
-------
factory_module : tvm.relay.backend.executor_factory.ExecutorFactoryModule
The runtime factory for the TVM graph executor.
"""
# pylint: enable=line-too-long
# fmt: on
if not isinstance(ir_mod, (IRModule, _function.Function)):
raise ValueError("Type of input parameter mod must be tvm.IRModule")
if isinstance(ir_mod, _function.Function):
if params:
ir_mod = bind_params_by_name(ir_mod, params)
ir_mod = IRModule.from_expr(ir_mod)
warnings.warn(
"Please use input parameter mod (tvm.IRModule) "
"instead of deprecated parameter mod (tvm.relay.function.Function)",
DeprecationWarning,
)
target = build_target_by_device_type_map(target)
if isinstance(target_host, (str, Target)):
target_host = Target(target_host)
elif target_host:
raise ValueError("target host must be the type of str, " +
"tvm.target.Target, or None")
target, target_host = Target.check_and_update_host_consist(
target, target_host, target_is_dict_key=False)
# Retrieve the executor from the target
executor = get_executor_from_target(target, target_host)
# If current dispatch context is fallback context (the default root context),
# then load pre-tuned parameters from TopHub
if isinstance(autotvm.DispatchContext.current, autotvm.FallbackContext):
tophub_context = autotvm.tophub.context(list(target.values()))
else:
tophub_context = autotvm.utils.EmptyContext()
with tophub_context:
bld_mod = BuildModule()
executor_config, runtime_mod, params = bld_mod.build(mod=ir_mod,
target=target,
params=params,
executor=executor,
mod_name=mod_name)
func_metadata = bld_mod.get_function_metadata()
if executor == "aot":
executor_factory = _executor_factory.AOTExecutorFactoryModule(
ir_mod, target, runtime_mod, mod_name, params, func_metadata)
elif executor == "graph":
executor_factory = _executor_factory.GraphExecutorFactoryModule(
ir_mod, target, executor_config, runtime_mod, mod_name, params,
func_metadata)
else:
assert False, "Executor " + executor + " not supported"
return executor_factory
def from_coreml(model, shape=None):
"""Convert from coreml model into Relay Function.
Parameters
----------
model:
coremltools.models.MLModel of a NeuralNetworkClassifier
shape : dict of str to int list/tuple, optional
The input shapes
Returns
-------
mod : tvm.IRModule
The relay module for compilation.
params : dict of str to tvm.nd.NDArray
The parameter dict to be used by Relay.
"""
try:
import coremltools as cm
except ImportError:
raise ImportError("The coremltools package must be installed")
assert isinstance(model, cm.models.MLModel)
spec = model.get_spec()
modeltype = spec.WhichOneof("Type")
assert modeltype in [
"neuralNetworkClassifier", "neuralNetwork", "neuralNetworkRegressor"
]
cc = getattr(spec, modeltype)
etab = ExprTable()
for i in spec.description.input:
input_shape = list(
shape[i.name]) if shape is not None and i.name in shape else None
etab.set_expr(i.name, _expr.var(i.name, shape=input_shape))
for pp in cc.preprocessing:
whichpp = pp.WhichOneof("preprocessor")
ppmethod = getattr(pp, whichpp)
if whichpp == "scaler":
# Be careful we maybe only preprocess one input when we have multi inputs
# which is stored in pp.featureName. See unit testing verify_image_scaler
# in test_forward.py for CoreML.
for i in spec.description.input:
# we have multi inputs
if len(spec.description.input) > 1:
assert pp.featureName != ""
if i.name == pp.featureName:
coreml_op_to_relay(ppmethod, i.name, i.name, etab)
else:
assert pp.featureName == ""
coreml_op_to_relay(ppmethod, i.name, i.name, etab)
else:
coreml_op_to_relay(ppmethod, pp.featureName, pp.featureName, etab)
for l in cc.layers:
layertype = l.WhichOneof("layer")
layerop = getattr(l, layertype)
if len(l.input) == 1:
coreml_op_to_relay(layerop, l.input[0], l.output, etab)
else:
coreml_op_to_relay(layerop, list(l.input), l.output, etab)
outexpr = [
etab.get_expr(o.name) if o.name in etab.exprs else _expr.var(o.name)
for o in spec.description.output
]
# check there are multiple outputs in the model and all are there in etab
multi_out = all(
[bool(o.name in etab.exprs) for o in spec.description.output])
outexpr = _expr.Tuple(outexpr) if multi_out else outexpr[0]
func = _function.Function(analysis.free_vars(outexpr), outexpr)
params = {
k: _nd.array(np.array(v, dtype=np.float32))
for k, v in etab.params.items()
}
return IRModule.from_expr(func), params
def extract_task_from_relay(
mod: Union[IRModule, RelayFunc],
target: Target,
params: Optional[Dict[str, NDArray]] = None,
*,
opt_level: int = 3,
pass_config: Optional[Dict[str, Any]] = None,
disabled_pass: Optional[List[str]] = None,
) -> List[ExtractedTask]:
"""Extract tuning tasks from a relay program.
Parameters
----------
mod : Union[tvm.IRModule, tvm.relay.Function]
The module or function to tune
target : tvm.target.Target
The compilation target
params : Optional[Dict[str, tvm.runtime.NDArray]]
The associated parameters of the program
opt_level : int
The optimization level of the compiler
pass_config : Optional[Dict[str, Any]]
The pass config of the compiler
disabled_pass : Optional[List[str]]
The list of disabled passes of the compiler
Returns
-------
tasks: List[ExtractedTask]
The tasks extracted from this network
"""
extract_task_func = get_global_func(
"relay.backend.MetaScheduleExtractTask")
assert extract_task_func
target = Target(target) if isinstance(target, str) else target
relay_params = {}
for name, param in params.items():
if isinstance(param, np.ndarray):
param = nd.array(param)
relay_params[name] = param
if disabled_pass is None:
disabled_pass = []
if pass_config is None:
pass_config = {"relay.backend.use_meta_schedule": True}
if isinstance(mod, RelayFunc):
mod = IRModule.from_expr(mod)
if not isinstance(target, Target):
target = Target(target)
with target, transform.PassContext(
opt_level=opt_level,
config=pass_config,
disabled_pass=disabled_pass,
):
tasks = extract_task_func(mod, target, relay_params)
# Tasks are extracted via post order visit, return the reversed list.
return list(reversed(tasks))
def from_keras(model, shape=None, layout='NCHW'):
"""Convert keras model to relay Function.
Parameters
----------
model : keras.engine.training.Model or tensorflow.keras.models.Model
The keras model to be converted.
shape: dict of str to int list/tuple
Input shapes of the model, optional
layout: str
One of 'NCHW' or 'NHWC', indicates how data should be arranged in
the output model. Default layout is 'NCHW' as it in general
performs better across TVM.
Returns
-------
mod : tvm.IRModule
The relay module for compilation.
params : dict of str to tvm.nd.NDArray
The parameter dict to be used by Relay.
"""
def _check_model_is_tf_keras():
return type(model).__module__.startswith("tensorflow.python.keras")
def _convert_input_layer(keras_layer):
input_name = keras_layer.name
input_shape = shape[
input_name] if shape is not None and input_name in shape else None
etab.set_expr(input_name, new_var(input_name, shape=input_shape))
is_tf_keras = _check_model_is_tf_keras()
if not is_tf_keras:
# Importing from Keras
try:
import keras
except ImportError:
raise ImportError("Keras must be installed")
if keras.backend.backend() != 'tensorflow':
raise ValueError(
"Keras frontend currently supports tensorflow backend only.")
if keras.backend.image_data_format() != 'channels_last':
raise ValueError(
"Keras frontend currently supports data_format = channels_last only."
)
expected_model_class = keras.engine.training.Model
input_layer_class = keras.engine.InputLayer
else:
# Importing from Tensorflow Keras (tf.keras)
try:
from tensorflow import keras as tf_keras
except ImportError:
raise ImportError("Tensorflow must be installed")
expected_model_class = tf_keras.models.Model
input_layer_class = tf_keras.layers.InputLayer
assert isinstance(model, expected_model_class)
etab = ExprTable()
# Set global data format.
assert layout in ['NCHW', 'NHWC',
'NDHWC'], "Layout must be one of 'NCHW', NHWC or NDHWC"
etab.data_layout = layout
for keras_layer in model.layers:
if isinstance(keras_layer, input_layer_class):
_convert_input_layer(keras_layer)
else:
inbound_nodes = keras_layer.inbound_nodes if hasattr(keras_layer, 'inbound_nodes') \
else keras_layer._inbound_nodes if hasattr(keras_layer, '_inbound_nodes') \
else None
if inbound_nodes is None:
raise TypeError(
"Unknown layer type or unsupported Keras version : {}".
format(keras_layer))
for node_idx, node in enumerate(inbound_nodes):
# If some nodes in imported model are not relevant to the current model,
# skip such layers.
# - In Keras, model._network_nodes contains keys of all nodes relevant to the
# current model;
# - In tf.Keras, this is already done as part of tensorflow.keras.network.get_config
if not is_tf_keras and \
not model._node_key(keras_layer, node_idx) in model._network_nodes:
continue
inexpr = []
# Since Keras allows creating multiple layers from the same name instance,
# we append node index to the expr name to make it unique.
# The one exception is InputLayer. Changing input variable names after conversion
# would confuse users, so we should keep them as far as possible. Fortunately,
# they are named uniquely to input_1, input_2, input_3... by default.
zip_node = zip(_as_list(node.node_indices),
_as_list(node.tensor_indices),
_as_list(node.inbound_layers))
for n_idx, t_idx, inbound_layer in zip_node:
if isinstance(inbound_layer, input_layer_class):
expr_name = inbound_layer.name
_convert_input_layer(inbound_layer)
else:
expr_name = inbound_layer.name + ':' + str(
n_idx) + ':' + str(t_idx)
expr = etab.get_expr(expr_name)
inexpr.append(expr)
if len(inexpr) == 1:
inexpr = inexpr[0]
keras_op_to_relay(inexpr, keras_layer,
keras_layer.name + ':' + str(node_idx), etab)
# model._output_coordinates contains out_node(oc[0]), node_index(oc[1]) and tensor_index(oc[2])
# Get all output nodes in etab using the name made from above values.
# The out exprs were added to etab in keras_op_to_relay using this name.
outexpr = [etab.get_expr(oc[0].name + ":" + str(oc[1]) + ":" + str(oc[2])) \
for oc in model._output_coordinates]
outexpr = outexpr[0] if len(outexpr) == 1 else _expr.Tuple(outexpr)
func = _function.Function(analysis.free_vars(outexpr), outexpr)
params = {
k: _nd.array(np.array(v, dtype=np.float32))
for k, v in etab.params.items()
}
return IRModule.from_expr(func), params
def extract_task_from_relay(
mod: IRModule,
target: Target,
params: Optional[Dict[str, NDArray]] = None,
*,
opt_level: int = 3,
pass_config: Optional[Dict[str, Any]] = None,
disabled_pass: Optional[List[str]] = None,
) -> List[ExtractedTask]:
"""Extract tuning tasks from a relay program.
Parameters
----------
mod : IRModule
The module or function to tune
target : tvm.target.Target
The compilation target
params : Optional[Dict[str, tvm.runtime.NDArray]]
The associated parameters of the program
opt_level : int
The optimization level of the compiler
pass_config : Optional[Dict[str, Any]]
The pass config of the compiler
disabled_pass : Optional[List[str]]
The list of disabled passes of the compiler
Returns
-------
tasks: List[ExtractedTask]
The tasks extracted from this network
"""
# pylint: disable=import-outside-toplevel
from tvm.relay import Function as RelayFunc
# pylint: enable=import-outside-toplevel
extract_task_func = get_global_func(
"relay.backend.MetaScheduleExtractTask",
allow_missing=False,
)
if isinstance(mod, RelayFunc):
mod = IRModule.from_expr(mod)
if not isinstance(target, Target):
target = Target(target)
if disabled_pass is None:
disabled_pass = []
if pass_config is None:
pass_config = {"relay.backend.use_meta_schedule": True}
if params is None:
params = {}
relay_params = {}
for name, param in params.items():
if isinstance(param, np.ndarray):
param = nd.array(param)
relay_params[name] = param
with autotvm_silencer(), target, transform.PassContext(
opt_level=opt_level,
config=pass_config,
disabled_pass=disabled_pass,
):
return list(extract_task_func(mod, target, relay_params))
def build(
ir_mod,
target=None,
target_host=None,
executor=Executor("graph"),
runtime=Runtime("cpp"),
workspace_memory_pools=None,
params=None,
mod_name="default",
):
# fmt: off
# pylint: disable=line-too-long
"""Helper function that builds a Relay function to run on TVM graph executor.
Parameters
----------
ir_mod : :py:class:`~tvm.IRModule`
The IR module to build. Using relay.Function is deprecated.
target : str, :any:`tvm.target.Target`, or dict of str(i.e. device/context name) to str/tvm.target.Target, optional
For heterogeneous compilation, it is a dictionary indicating context to
target mapping. For homogeneous compilation, it is a build target.
target_host : str or :any:`tvm.target.Target`, optional
Host compilation target, if target is device.
When TVM compiles device specific program such as CUDA,
we also need host(CPU) side code to interact with the driver
setup the dimensions and parameters correctly.
target_host is used to specify the host side codegen target.
By default, llvm is used if it is enabled,
otherwise a stackvm interpreter is used.
executor : Optional[Executor]
The executor configuration with which to build the model.
Defaults to "graph" if no executor specified.
runtime : Optional[Runtime]
Runtime configuration to use when building the model.
Defaults to "cpp" if no runtime specified.
workspace_memory_pools : Optional[WorkspaceMemoryPools]
The object that contains an Array of PoolInfo objects
that hold properties of workspace pools that could be
used by the inference.
params : dict of str to NDArray
Input parameters to the graph that do not change
during inference time. Used for constant folding.
mod_name: Optional[str]
The module name we will build
Returns
-------
factory_module : tvm.relay.backend.executor_factory.ExecutorFactoryModule
The runtime factory for the TVM graph executor.
"""
# pylint: enable=line-too-long
# fmt: on
if not isinstance(ir_mod, (IRModule, _function.Function)):
raise ValueError("Type of input parameter mod must be tvm.IRModule")
if isinstance(ir_mod, _function.Function):
if params:
ir_mod = bind_params_by_name(ir_mod, params)
ir_mod = IRModule.from_expr(ir_mod)
warnings.warn(
"Please use input parameter mod (tvm.IRModule) "
"instead of deprecated parameter mod (tvm.relay.function.Function)",
DeprecationWarning,
)
if target_host is not None:
warnings.warn(
"target_host parameter is going to be deprecated. "
"Please pass in tvm.target.Target(target, host=target_host) instead."
)
target, target_host = Target.check_and_update_host_consist(
target, target_host, target_is_dict_key=False)
target = build_target_by_device_type_map(target)
if isinstance(target_host, (str, Target)):
target_host = Target(target_host)
elif target_host:
raise ValueError("target host must be the type of str, " +
"tvm.target.Target, or None")
# All of this logic is to raise deprecation warnings for various parameters
# TODO(Mousius) Remove these after some time
deprecated_params_target = target_host or list(target.values())[0]
deprecated_executor, deprecated_runtime = _reconstruct_from_deprecated_options(
deprecated_params_target)
if deprecated_executor:
executor = deprecated_executor
if deprecated_runtime:
runtime = deprecated_runtime
# If current dispatch context is fallback context (the default root context),
# then load pre-tuned parameters from TopHub
if isinstance(autotvm.DispatchContext.current, autotvm.FallbackContext):
tophub_context = autotvm.tophub.context(list(target.values()))
else:
tophub_context = autotvm.utils.EmptyContext()
with tophub_context:
bld_mod = BuildModule()
graph_json, runtime_mod, params = bld_mod.build(
mod=ir_mod,
target=target,
params=params,
executor=executor,
runtime=runtime,
workspace_memory_pools=workspace_memory_pools,
mod_name=mod_name,
)
func_metadata = bld_mod.get_function_metadata()
devices = bld_mod.get_devices()
lowered_ir_mods = bld_mod.get_irmodule()
executor_codegen_metadata = bld_mod.get_executor_codegen_metadata()
if str(executor) == "aot":
executor_factory = _executor_factory.AOTExecutorFactoryModule(
ir_mod,
lowered_ir_mods,
target,
executor,
runtime,
runtime_mod,
mod_name,
params,
func_metadata,
executor_codegen_metadata,
devices,
)
elif str(executor) == "graph":
executor_factory = _executor_factory.GraphExecutorFactoryModule(
ir_mod, target, executor, graph_json, runtime_mod, mod_name,
params, func_metadata)
else:
assert False, "Executor " + executor + " not supported"
return executor_factory
def test_partition():
in_1 = relay.var("in_1", shape=(10, 10), dtype="float32")
in_2 = relay.var("in_2", shape=(10, 10), dtype="float32")
in_3 = relay.var("in_3", shape=(10, 10), dtype="float32")
in_4 = relay.var("in_4", shape=(10, 10), dtype="float32")
in_5 = relay.var("in_5", shape=(10, 10), dtype="float32")
in_6 = relay.var("in_6", shape=(10, 10), dtype="float32")
in_7 = relay.var("in_7", shape=(10, 10), dtype="float32")
in_8 = relay.var("in_8", shape=(10, 10), dtype="float32")
in_9 = relay.var("in_9", shape=(10, 10), dtype="float32")
in_10 = relay.var("in_10", shape=(10, 10), dtype="float32")
begin0 = compiler_begin(in_1, "onnx")
begin1 = compiler_begin(in_2, "onnx")
begin2 = compiler_begin(in_3, "onnx")
begin3 = compiler_begin(in_4, "onnx")
node0 = relay.add(begin0, begin1)
node1 = relay.add(begin2, begin3)
end0 = compiler_end(node0, "onnx")
end1 = compiler_end(node1, "onnx")
begin4 = compiler_begin(end0, "onnx")
begin5 = compiler_begin(end1, "onnx")
node2 = relay.add(begin4, begin5)
end2 = compiler_end(node2, "onnx")
dbegin0 = compiler_begin(in_5, "default")
dbegin1 = compiler_begin(in_6, "default")
node3 = relay.subtract(dbegin0, dbegin1)
dbegin2 = compiler_begin(in_7, "default")
dend1 = compiler_end(node3, "default")
dbegin3 = compiler_begin(dend1, "default")
node4 = relay.subtract(dbegin2, dbegin3)
dend2 = compiler_end(node4, "default")
begin6 = compiler_begin(end2, "onnx")
begin7 = compiler_begin(dend2, "onnx")
node5 = relay.add(begin6, begin7)
end3 = compiler_end(node5, "onnx")
end4 = compiler_end(node5, "onnx")
dbegin4 = compiler_begin(in_8, "default")
dbegin5 = compiler_begin(end3, "default")
node6 = relay.subtract(dbegin4, dbegin5)
begin8 = compiler_begin(in_9, "onnx")
begin9 = compiler_begin(end4, "onnx")
node7 = relay.multiply(begin8, begin9)
end5 = compiler_end(node7, "onnx")
dend3 = compiler_end(node6, "default")
begin10 = compiler_begin(dend3, "onnx")
begin11 = compiler_begin(end5, "onnx")
node8 = relay.add(begin10, begin11)
end6 = compiler_end(node8, "onnx")
begin12 = compiler_begin(in_10, "onnx")
begin13 = compiler_begin(end6, "onnx")
node9 = relay.add(begin12, begin13)
end7 = compiler_end(node9, "onnx")
func = relay.Function(
[in_1, in_2, in_3, in_4, in_5, in_6, in_7, in_8, in_9, in_10], end7)
target = "llvm"
mod = IRModule.from_expr(func)
mod = transform.PartitionGraph()(mod)
with tvm.transform.PassContext(opt_level=3, disabled_pass=["FuseOps"]):
graph_json, mod1, params = relay.build(mod, target)
assert mod1.type_key == "metadata"
assert mod1.imported_modules[0].type_key == "llvm"
assert mod1.imported_modules[0].get_source()
assert mod1.imported_modules[1].type_key == "onnx"
assert mod1.imported_modules[1].get_source()
def from_caffe(init_net, predict_net, shape_dict, dtype_dict):
"""Convert from caffe model into compatible relay Function.
Parameters
----------
init_net : caffe_pb2.NetParameter
caffemodel
predict_net : caffe_pb2.NetParameter
caffe prototxt
shape_dict : dict of str to int list/tuple
Input shapes of the model.
dtype_dict : dict of str to str
Input types of the model.
Returns
-------
mod : tvm.IRModule
The relay module for compilation.
params : dict of str to tvm.NDArray
The parameter dict to be used by relay
"""
old_caffe = False
if len(predict_net.input) != 0: # old caffe version
old_caffe = True
model_inputs = list(predict_net.input)
predict_layer = predict_net.layer
# replace layer's top with its name and update other layers'bottoms
_rebuild_layers(predict_layer)
# obtain inputs and outputs of Net
if old_caffe:
_, model_outputs = _get_inputs_outputs(predict_layer)
else:
model_inputs, model_outputs = _get_inputs_outputs(predict_layer)
exp_tab = ExprTable()
for in_name in model_inputs:
shape = shape_dict[in_name] if in_name in shape_dict else None
dtype = dtype_dict[in_name] if in_name in dtype_dict else "float32"
exp_tab.set_expr(in_name, _expr.var(in_name, shape=shape, dtype=dtype))
if list(init_net.layer):
init_layer = init_net.layer
else:
init_layer = init_net.layers
init_layer_dict = {il.name: il for il in init_layer}
# op code in model
op_converter = OperatorConverter(init_layer_dict, predict_layer, exp_tab)
op_converter.check_unsupported_ops()
op_converter.op_fuse()
op_converter.convert_op_to_relay()
# params and outputs
params = {k: _nd.array(np.array(v)) for k, v in exp_tab.params.items()}
outputs = list()
for n in model_outputs:
if n in op_converter.changed_layers:
n = op_converter.changed_layers[n]
outputs.append(exp_tab.get_expr(n))
outputs = outputs[0] if len(outputs) == 1 else _expr.Tuple(outputs)
func = _function.Function(analysis.free_vars(outputs), outputs)
mod = IRModule.from_expr(func)
return mod, params
def extract_task_from_relay(
mod: IRModule,
target: Target,
params: Optional[Dict[str, NDArray]] = None,
*,
opt_level: int = 3,
pass_config: Optional[Dict[str, Any]] = None,
disabled_pass: Optional[List[str]] = None,
te_filter_func: Union[str, None, Callable[[List[Tensor]],
PrimFunc]] = None,
) -> List[ExtractedTask]:
"""Extract tuning tasks from a relay program.
Parameters
----------
mod : IRModule
The module or function to tune
target : tvm.target.Target
The compilation target
params : Optional[Dict[str, tvm.runtime.NDArray]]
The associated parameters of the program
opt_level : int
The optimization level of the compiler
pass_config : Optional[Dict[str, Any]]
The pass config of the compiler
disabled_pass : Optional[List[str]]
The list of disabled passes of the compiler
te_filter_func : Callable[[List[tvm.te.Tensor]], bool]
The filter function to filter out the extracted tasks
If it's a string, it's the name of the filtering function. Built in functions are
- "meta_schedule.DefaultTaskFilter"
- "meta_schedule.DefaultTaskFilterAllowExtern"
If it's None, it's the default filtering function
If it's a callable, it's the filtering function
Returns
-------
tasks: List[ExtractedTask]
The tasks extracted from this network
"""
# pylint: disable=import-outside-toplevel
from tvm import autotvm
from tvm.relay import Function as RelayFunc
# pylint: enable=import-outside-toplevel
if isinstance(te_filter_func, str):
te_filter_func = get_global_func(te_filter_func)
extract_task_func = get_global_func(
"relay.backend.MetaScheduleExtractTask",
allow_missing=False,
)
if isinstance(mod, RelayFunc):
mod = IRModule.from_expr(mod)
if not isinstance(target, Target):
target = Target(target)
if disabled_pass is None:
disabled_pass = []
if pass_config is None:
pass_config = {"relay.backend.use_meta_schedule": True}
if params is None:
params = {}
relay_params = {}
for name, param in params.items():
if isinstance(param, np.ndarray):
param = nd.array(param)
relay_params[name] = param
with target, autotvm_silencer(), transform.PassContext(
opt_level=opt_level,
config=pass_config,
disabled_pass=disabled_pass,
):
if target.kind.name != "cuda" and isinstance(
autotvm.DispatchContext.current, autotvm.FallbackContext):
tophub_context = autotvm.tophub.context(target)
else:
tophub_context = autotvm.utils.EmptyContext()
with tophub_context:
return list(
extract_task_func(mod, target, relay_params, te_filter_func))
def build(mod, target=None, target_host=None, params=None, mod_name='default'):
"""Helper function that builds a Relay function to run on TVM graph
runtime.
Parameters
----------
mod : :py:class:`~tvm.IRModule`
The IR module to build. Using relay.Function is deprecated.
target : str, :any:`tvm.target.Target`, or dict of str(i.e. device/context
name) to str/tvm.target.Target, optional
For heterogeneous compilation, it is a dictionary indicating context to
target mapping. For homogeneous compilation, it is a build target.
target_host : str or :any:`tvm.target.Target`, optional
Host compilation target, if target is device.
When TVM compiles device specific program such as CUDA,
we also need host(CPU) side code to interact with the driver
setup the dimensions and parameters correctly.
target_host is used to specify the host side codegen target.
By default, llvm is used if it is enabled,
otherwise a stackvm intepreter is used.
params : dict of str to NDArray
Input parameters to the graph that do not change
during inference time. Used for constant folding.
mod_name: Optional[str]
The module name we will build
Returns
-------
graph_json : str
The json string that can be accepted by graph runtime.
mod : tvm.Module
The module containing necessary libraries.
params : dict
The parameters of the final graph.
"""
if not isinstance(mod, (IRModule, _function.Function)):
raise ValueError("Type of input parameter mod must be tvm.IRModule")
if isinstance(mod, _function.Function):
if params:
mod = bind_params_by_name(mod, params)
mod = IRModule.from_expr(mod)
warnings.warn(
"Please use input parameter mod (tvm.IRModule) "
"instead of deprecated parameter mod (tvm.relay.function.Function)",
DeprecationWarning)
target = _update_target(target)
if isinstance(target_host, (str, _target.Target)):
target_host = _target.create(target_host)
elif target_host:
raise ValueError("target host must be the type of str, " +
"tvm.target.Target, or None")
# If current dispatch context is fallback context (the default root context),
# then load pre-tuned parameters from TopHub
if isinstance(autotvm.DispatchContext.current, autotvm.FallbackContext):
tophub_context = autotvm.tophub.context(list(target.values()))
else:
tophub_context = autotvm.util.EmptyContext()
with tophub_context:
bld_mod = BuildModule()
graph_json, mod, params = bld_mod.build(mod, target, target_host,
params)
mod = _graph_runtime_factory.GraphRuntimeFactoryModule(
graph_json, mod, mod_name, params)
return mod
def extract_task_from_relay(
mod: Union[IRModule, RelayFunc],
target: Target,
params: Optional[Dict[str, NDArray]] = None,
*,
opt_level: int = 3,
pass_config: Dict[str, Any] = {
"relay.backend.use_meta_schedule": True,
},
disabled_pass: List[str] = [],
) -> List[ExtractedTask]:
"""Extract tuning tasks from a relay program.
Parameters
----------
mod : Union[tvm.IRModule, tvm.relay.Function]
The module or function to tune
target : tvm.target.Target
The compilation target
params : Optional[Dict[str, tvm.runtime.NDArray]]
The associated parameters of the program
opt_level : int
The optimization level of the compiler
pass_config : Dict[str, Any]
The pass config of the compiler
disabled_pass : List[str]
The list of disabled passes of the compiler
Returns
-------
tasks: List[ExtractedTask]
The tasks extracted from this network
"""
@contextmanager
def _autotvm_silencer():
from tvm import autotvm # pylint: disable=import-outside-toplevel
silent = autotvm.GLOBAL_SCOPE.silent
autotvm.GLOBAL_SCOPE.silent = True
try:
yield
finally:
autotvm.GLOBAL_SCOPE.silent = silent
def _thread_run(func: Callable[[], None]) -> None:
import threading # pylint: disable=import-outside-toplevel
thread = threading.Thread(target=func)
thread.start()
thread.join()
env = TaskExtraction()
if isinstance(mod, RelayFunc):
mod = IRModule.from_expr(mod)
if not isinstance(target, Target):
target = Target(target)
def _func():
with env, _autotvm_silencer(), transform.PassContext(
config=pass_config,
disabled_pass=disabled_pass,
opt_level=opt_level,
):
compiler = vm.VMCompiler()
if params:
compiler.set_params(params)
compiler.lower(mod, target)
_thread_run(_func)
return env.tasks
