def check_model(sym, shapes, dtypes, params): net = nnvm.graph.create(sym) graph_json, mod, params = nnvm.compiler.build(net, 'llvm', shape=shapes, dtype=dtypes, params=params) nnvm_rts = graph_runtime.create(graph_json, mod, tvm.cpu(0)) inputs = {} for name in shapes: np_array = np.random.rand(*shapes[name]).astype('float32') inputs[name] = tvm.nd.array(np_array) nnvm_rts.set_input(**params) nnvm_rts.run(**inputs) nnvm_out = nnvm_rts.get_output(0) relay_model, params = to_relay.to_relay(net, shapes, dtypes, params) mod = tvm.relay.Module.from_expr(relay_model) mod = transform.InferType()(mod) relay_rts = create_executor(kind='graph', mod=mod, ctx=tvm.cpu(0), target='llvm') inputs.update(params) relay_out = relay_rts.evaluate()(*list(inputs.values())) np.testing.assert_allclose(nnvm_out.asnumpy(), relay_out.asnumpy())
def get_relay_output(sym, x, params, target, ctx, dtype='float32'): shape_dict = {'data': x.shape} func, params = to_relay(sym, shape_dict, dtype, params) with relay.build_config(opt_level=3): graph, lib, params = relay.build(func, target=target, params=params) m = graph_runtime.create(graph, lib, ctx) m.set_input("data", tvm.nd.array(x.astype(dtype))) m.set_input(**params) m.run() return m.get_output(0).asnumpy()
def get_tvm_output(xs, target, ctx, dtype='float32'): sym, params = nnvm.frontend.from_keras(keras_model) shape_dict = { name: x.shape for (name, x) in zip(keras_model.input_names, xs) } with relay.build_module.build_config(opt_level=2): func, params = to_relay.to_relay(sym, shape_dict, dtype, params) graph, lib, params = relay.build(func, target='llvm', params=params) m = graph_runtime.create(graph, lib, ctx) for name, x in zip(keras_model.input_names, xs): m.set_input(name, tvm.nd.array(x.astype(dtype))) m.set_input(**params) m.run() return [m.get_output(i).asnumpy() for i in range(m.get_num_outputs())]
def check_function(symbol, forward=None, backward=None, grad_input_vars=None, shape=None, dtype=None, in_range=None, values=None, exclude_targets=None, only_targets=None, additional_params=None, numerical_grads=None, numerical_grads_params=None, atol=1e-5, rtol=1e-5, quiet=False): """Compute the function and/or its gradients on a random input and raise an exception if the result doesn't match the reference implementation. Parameters ---------- symbol : nnvm.Symbol A symbol representing the output. forward : Callable[..., List[numpy.ndarray]], optional A reference implementation to compare with. backward : Callable[..., List[numpy.ndarray] or Dict[str, numpy.ndarray]], optional A reference implementation of gradients. Should also accept head_grads besides normal inputs which is a list of gradients of some scalar wrt the outputs or just a single gradient if there are multiple outputs. Should return either a dict mapping input variable names to the respective gradients or a list of gradients wrt variables from grad_input_vars in exactly the same order (in alphabetical order by default). grad_input_vars : List[nnvm.Symbol or str], optional A list of variables with respect to which the gradients will be computed. None (default) means that all input variables will be used in an alphabetical order. shape : Dict[nnvm.Symbol or str, Tuple[int]] or Tuple[int], optional A dict mapping input variable names to shapes, or just a single shape. By default shapes will be inferred from variables' attributes (see the Examples). Note that this parameter takes precedence over variables' attributes. dtype : Dict[nnvm.Symbol or str, str] or str, optional A dict mapping input variable names to dtypes, or just a single dtype. By default dtypes will be inferred from variables' attributes (see the Examples). If dtypes cannot be inferred for some variables then float32 will be used as a fallback. Note that this parameter takes precedence over variables' attributes. in_range : Dict[nnvm.Symbol or str, (float, float)] or (float, float), optional A dict mapping input variable names to ranges or just a single range (the same for all variables). Input values will be generated from uniform distributions on these ranges. `head_grads` can also be assigned a range this way. values : Dict[nnvm.Symbol or str, numpy.ndarray], optional A dict explicitly providing values for some variables instead of random generation. exclude_targets : Set[str], optional Skip compiling and running anything for these targets. only_targets : Set[str], optional Test only for those targets from `ctx_list()` that are also in this set. additional_params : dict, optional A dict of additional parameters which will be passed to forward and backward. numerical_grads : bool or 'if_possible', optional Whether to additionally check against numerically computed gradients. If 'if_possible' or None is passed (which is the default) then it will try to create a gradient computation graph and then check gradients numerically only if this graph can be created (i.e. if there are some operations with unimplemented gradients, it will just issue a warning). Checking against numerical gradients is done via the `check_numerical_grads` function. numerical_grads_params : dict, optional Additional parameters for `check_numerical_grads`. atol : float, optional Absolute tolerance for `tvm.testing.assert_allclose`. NOT used for numerical gradients. rtol : float, optional Relative tolerance for `tvm.testing.assert_allclose`. NOT used for numerical gradients. quiet : bool, optional Don't dump additional information to stdout on failure. Examples -------- .. code-block:: python x = sym.Variable("x", shape=(1, 2)) y = sym.Variable("y", shape=(1, 2)) # check the function and its gradients both numerically and using a reference function check_function(x + 2*y, lambda x, y: x + 2*y, lambda x, y, head_grads: {'x': head_grads, 'y': 2*head_grads}) # just check gradients numerically check_function(x + 2*y, numerical_grads=True) # just check the forward computation check_function(x + 2*y, lambda x, y: x + 2*y, numerical_grads=False) # specifying dtype check_function(x + 2*y, lambda x, y: x + 2*y, dtype='float64') # dtypes can also be specified during variable creation with dtype codes x = sym.Variable("x", dtype=0) check_function(x + 1, shape=(2, 2), numerical_grads=True) """ # validate and preprocess the input params if numerical_grads is None and forward is None and backward is None: raise ValueError( "No reference function was passed to check_function. If you only want to " "check gradients numerically, pass numerical_grads=True explicitly." ) if numerical_grads is None: numerical_grads = 'if_possible' if numerical_grads not in [False, True, 'if_possible']: raise ValueError( "numerical_grads must be a bool or 'if_possible', not {}".format( numerical_grads)) if additional_params is None: additional_params = {} input_vars = symbol.list_input_variables() input_dict = {x.attr('name'): x for x in input_vars} if grad_input_vars is None: grad_input_vars = sorted(input_vars, key=lambda x: x.attr('name')) else: grad_input_vars = [ input_dict[x] if isinstance(x, str) else x for x in grad_input_vars ] in_range = _dict_var_to_dict_str(in_range) values = _dict_var_to_dict_str(values) out_len = len(symbol.list_output_names()) # Infer the output shapes and dtypes, and preprocess the shape and dtype params forward_graph, shape, dtype, out_shapes, out_dtypes = \ infer_shapes_dtypes(nnvm.graph.create(symbol), shape=shape, dtype=dtype, fallback_dtype='float32') if not all(out_shapes) or not all(out_dtypes): if not quiet: print(forward_graph.ir(join_node_attrs=['shape', 'dtype'])) raise ValueError("Could not infer shapes or dtypes for outputs.\n" "out_shapes = {}\nout_dtypes = {}".format( out_shapes, out_dtypes)) backward_graph = None # If we want gradients, we have to recreate the graph, but now with gradient computations # Note that here we need out_shapes for defining the shape of head grads, so we have to # create the graph twice if backward is not None or numerical_grads: try: head_grads_symbols = [ nnvm.symbol.Variable("head_grads_" + str(i), shape=out_shapes[i], dtype=DTYPE_TO_TCODE[out_dtypes[i]]) for i in range(out_len) ] grad_symbols = graph_util.gradients([symbol], grad_input_vars, grad_ys=head_grads_symbols) # Sometimes grads do not depend on head_grads, so head_grads does not appear # in the variable list; adding it manually prevents this, making things a bit easier backward_graph = \ nnvm.graph.create(nnvm.symbol.Group([symbol] + grad_symbols + head_grads_symbols)) backward_graph, shape, dtype, out_shapes, out_dtypes = \ infer_shapes_dtypes(backward_graph, shape=shape, dtype=dtype, fallback_dtype='float32') except nnvm._base.NNVMError as err: if backward is None and numerical_grads == "if_possible": logging.warning("Won't check gradients because: %s", str(err).split('\n', 1)[0]) numerical_grads = False backward_graph = None else: raise main_graph = backward_graph if backward_graph is not None else forward_graph # Generate random data for inputs (including head_grads) np_inputs = {} for x in main_graph.symbol.list_input_variables(): x_name = x.attr('name') x_shape = shape[x_name] x_dtype = dtype[x_name] if values is not None and x_name in values: np_inputs[x_name] = values[x_name].astype(x_dtype) continue low = -1.0 high = 1.0 if in_range is not None: if isinstance(in_range, dict): if x_name in in_range: low = in_range[x_name][0] high = in_range[x_name][1] else: low = in_range[0] high = in_range[1] np_inputs[x_name] = np.random.uniform(size=x_shape, low=low, high=high).astype(x_dtype) np_inputs_without_head_grads = { k: np_inputs[k] for k in np_inputs if not k.startswith('head_grads_') } nothing_was_done = True # Compute and compare the results for target, ctx in ctx_list(): if exclude_targets is not None: if target in exclude_targets or str(target) in exclude_targets: logging.info("Skipping target = %s, ctx = %s", target, ctx) continue if only_targets is not None: if target not in only_targets and str(target) not in only_targets: logging.info("Skipping target = %s, ctx = %s", target, ctx) continue logging.info("Checking computation on target = %s, ctx = %s", target, ctx) debug_stage = None try: nnvm_res = None debug_stage = "compiling" main_function = graph_to_function(main_graph, target, ctx) # nnvm_res contains the output and gradients (if they are needed) debug_stage = "running" nnvm_res = main_function(**np_inputs) try: logging.debug("checking to_relay conversion") inputs = np_inputs_without_head_grads.copy() func, inputs = to_relay(main_graph, shape, dtype, params=inputs) with relay.build_config(opt_level=3): graph, lib, params = relay.build(func, target=target) m = graph_runtime.create(graph, lib, ctx) m.set_input(**inputs) m.set_input(**params) m.run() for i in range(out_len): relay_out = m.get_output(i).asnumpy() tvm.testing.assert_allclose(nnvm_res[i], relay_out, atol=atol, rtol=rtol) except NotImplementedError as err: # the NNVM operator is not supported yet logging.warning(err) if backward_graph is not None: grad_var_names = [x.attr('name') for x in grad_input_vars] nnvm_grads = { x: v for x, v in zip(grad_var_names, nnvm_res[out_len:]) } if forward is not None: nothing_was_done = False debug_stage = "checking forward computation" logging.debug(debug_stage) params = {} params.update(np_inputs_without_head_grads) params.update(additional_params) numpy_res = forward(**params) if isinstance(numpy_res, tuple): numpy_res = list(numpy_res) if not isinstance(numpy_res, list): numpy_res = [numpy_res] if len(numpy_res) != out_len: raise ValueError( "Forward function returned {} values, but " "the nnvm graph returns {} values".format( len(numpy_res), out_len)) for i in range(out_len): tvm.testing.assert_allclose(nnvm_res[i], numpy_res[i], atol=atol, rtol=rtol) if backward is not None: nothing_was_done = False debug_stage = "checking gradients" logging.debug(debug_stage) np_head_grads = [ np_inputs["head_grads_" + str(i)] for i in range(out_len) ] if out_len == 1: np_head_grads = np_head_grads[0] params = {'head_grads': np_head_grads} params.update(np_inputs_without_head_grads) params.update(additional_params) numpy_grads = backward(**params) if not isinstance(numpy_grads, dict): if isinstance(numpy_grads, tuple): numpy_grads = list(numpy_grads) if not isinstance(numpy_grads, list): numpy_grads = [numpy_grads] numpy_grads = { x: v for x, v in zip(grad_var_names, numpy_grads) } if len(numpy_grads) != len(grad_var_names): raise ValueError( "The backward function returns a list of gradients which " "does not contain gradients for these variables: {}" .format(set(grad_var_names) - set(numpy_grads))) for x_name in numpy_grads: tvm.testing.assert_allclose(nnvm_grads[x_name], numpy_grads[x_name], atol=atol, rtol=rtol) if numerical_grads: nothing_was_done = False debug_stage = "checking gradients numerically" logging.debug(debug_stage) forward_function = graph_to_function(forward_graph, target, ctx) # Since the result may be non-scalar, we have to put another operation on the top, # so we just multiple by the randomly generated head_grads and then sum everything. # This way we can reuse the gradient values which has been already computed. def scalar_function(**kwargs): res = forward_function(**kwargs) return np.sum([ np.dot(np_inputs['head_grads_' + str(i)].ravel(), res[i].ravel()) for i in range(out_len) ]) if numerical_grads_params is None: numerical_grads_params = {} check_numerical_grads( scalar_function, input_values=np_inputs_without_head_grads, grad_values=nnvm_grads, **numerical_grads_params) except: if not quiet: print( "\ncheck_function failed while {}, here is the main graph". format(debug_stage)) print(main_graph.ir(join_node_attrs=['shape', 'dtype'])) if nnvm_res is not None: print("Generated inputs:") print(np_inputs) print() raise if nothing_was_done: logging.warning( "Nothing was done in check_function. Check ctx_list().")
if local: target = 'llvm' ctx = tvm.cpu(0) else: target = 'cuda' ctx = tvm.gpu(0) data = np.empty([batch_size, net.c, net.h, net.w], dtype) shape = {'data': data.shape} dtype_dict = {} # convert nnvm to relay print("convert nnvm symbols into relay function...") from nnvm.to_relay import to_relay func, params = to_relay(sym, shape, 'float32', params=params) # optimization print("optimize relay graph...") with tvm.relay.build_config(opt_level=2): func = tvm.relay.optimize(func, target, params) # quantize print("apply quantization...") from tvm.relay import quantize with quantize.qconfig(): func = quantize.quantize(func, params) # Relay build print("Compiling the model...") print(func.astext(show_meta_data=False)) with tvm.relay.build_config(opt_level=3): graph, lib, params = tvm.relay.build(func, target=target, params=params)