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