nnvm.compiler

NNVM compiler toolchain.

User only need to use build and build_config to do the compilation, and save_param_dict to save the parameters into bytes. The other APIs are for more advanced interaction with the compiler toolchain.

nnvm.compiler.build(graph, target=None, shape=None, dtype=’float32’, params=None, target_host=None)

Build graph into runtime library.

The build function will optimize the graph and do the compilation.

When params is provided, the compiler might split the graph to pre-compute certain values, so the final execution graph can be different from the original one.

Parameters:
  • graph (Graph) – The graph to be used in lowering
  • target (str or tvm.target.Target, optional) – The build target
  • shape (dict of str to tuple, optional) – The input shape to the graph
  • dtype (str or dict of str to str) – The input types to the graph
  • params (dict of str to NDArray) – Input parameters to the graph that do not change during inference time. Used for pre-compute folding optimization.
  • target_host (str or 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.
Returns:

  • graph (Graph) – The final execution graph.
  • libmod (tvm.Module) – The modue that comes with the execution graph
  • params (dict of str to NDArray) – The updated parameters of graph if params is passed. This can be different from the params passed in.

nnvm.compiler.build_config(**kwargs)

Configure the build behavior by setting config variables.

Parameters:
  • opt_level (int, default=2) – Optimization level. See OPT_PASS_LEVEL for level of each pass.
  • add_pass (set of str) – Optimization pass to be added regardless of optimization level.
Returns:

config – The build configuration

Return type:

BuildConfig

nnvm.compiler.save_param_dict(params)

Save parameter dictionary to binary bytes.

The result binary bytes can be loaded by the GraphModule with API “load_params”.

Parameters:params (dict of str to NDArray) – The parameter dictionary.
Returns:param_bytes – Serialized parameters.
Return type:bytearray

Examples

# compile and save the modules to file.
graph, lib, params = nnvm.compiler.build(
   graph, target, shape={"data", data_shape}, params=params)
module = graph_runtime.create(graph, lib, tvm.gpu(0))
# save the parameters as byte array
param_bytes = nnvm.compiler.save_param_dict(params)
# We can serialize the param_bytes and load it back later.
# Pass in byte array to module to directly set parameters
module["load_params"](param_bytes)
nnvm.compiler.load_param_dict(param_bytes)

Load parameter dictionary to binary bytes.

Parameters:param_bytes (bytearray) – Serialized parameters.
Returns:params – The parameter dictionary.
Return type:dict of str to NDArray
nnvm.compiler.optimize(graph, shape, dtype=’float32’)

Perform target and parameter invariant graph optimization.

This is an advanced function that usually do not need to be called. Call build instead.

Parameters:graph (Graph) – The graph to be used in optimized.
Returns:graph – The optimized graph.
Return type:Graph

Utility function to get information from graph.

nnvm.compiler.graph_util.check_graph_equal(grapha, graphb, compare_variable_attrs=False)

Check if two graphs have equal structure.

Parameters:
  • grapha (Graph) – The first graph
  • graphb (Graph) – The second graph
  • compare_variable_attrs (bool, optional) – Whether we want to compare attributes(names) on variables. Usually it is safe to skip it unless we want input name to exactly match
Raises:

ValueError – ValueError is raised with error message when graph not equal

nnvm.compiler.graph_util.get_gradient_graph(ys, xs, grad_ys=None)

Create gradient graph of ys with respect to xs.

Parameters:
  • ys (Symbol or list of Symbol) – Symbols from which the gradient is calculated.
  • xs (Symbol or list of Symbol) – Symbols the gradient respect to. For group symbol, gradients for all outputs will be calculated.
  • grad_ys (Symbol or list of Symbol) – Head gradients for ys.
Returns:

ret – Generated gradient graph.

Return type:

Graph

nnvm.compiler.graph_util.gradients(ys, xs, grad_ys=None)

Create gradient symbol of ys respect to xs.

Parameters:
  • ys (Symbol or list of Symbol) – Symbols from which the gradient is calculated.
  • xs (Symbol or list of Symbol) – Symbols the gradient respect to. For group symbol, gradients for all outputs will be calculated.
  • grad_ys (Symbol or list of Symbol) – Head gradients for ys.
Returns:

ret – Generated gradient symbol. For each xs, all gradients from ys are merged into a single symbol.

Return type:

list of Symbol

nnvm.compiler.graph_util.infer_dtype(graph, **dtype)

Infer the type given the typeS of inputs.

Parameters:
  • graph (Graph) – The graph to perform type inference from
  • dtype (dict of str to dtype) – The specific input data type.
Returns:

  • in_dtype (list of tuple) – Dtype of inputs
  • out_dtype (list of tuple) – Dtype of outputs

nnvm.compiler.graph_util.infer_shape(graph, **shape)

Infer the shape given the shape of inputs.

Parameters:
  • graph (Graph) – The graph to perform shape inference from
  • shape (dict of str to tuple) – The specific input shape.
Returns:

  • in_shape (list of tuple) – Shape of inputs
  • out_shape (list of tuple) – Shape of outputs

Utilities to access graph attributes

nnvm.compiler.graph_attr.set_dtype_inputs(g, dtype)

Set the dtype inputs of graph nodes

Parameters:
  • g (Graph) – The input graph
  • dtype (dict of str to str or str) – The input dtype
Returns:

g – The updated graph with updated dtype.

Return type:

Graph

nnvm.compiler.graph_attr.set_layout_inputs(g, layout)

Set the layout inputs of graph nodes

Parameters:
  • g (Graph) – The input graph
  • layout (dict of str to str or str) – The input layout
Returns:

g – The updated graph with updated dtype.

Return type:

Graph

nnvm.compiler.graph_attr.set_shape_inputs(g, shape)

Set the shape of input graph nodes in the graph attribute.

Parameters:
  • g (Graph) – The input graph
  • shape (dict of str to tuple) – The input shape
Returns:

g – The updated graph with updated shape.

Return type:

Graph

Compiler engine interface to internal engine

You can get the engine singleton at nnvm.compiler.engine

class nnvm.compiler.compile_engine.Engine

Global singleton compilation engine.

You can get the singleton at nnvm.compiler.engine

clear_cache()

Clear the existing cached functions.

dump()

Return a string representation of engine dump

Returns:dump – The dumped string representation
Return type:str
items()

List the available cache key value pairs.

Returns:item_list – The existing cache items
Return type:list of (GraphKey, GraphCacheEntry)
class nnvm.compiler.compile_engine.GraphCacheEntry(handle)

CacheEntry of compilation into a TVM Function

class nnvm.compiler.compile_engine.GraphFunc(handle)

Compiled result of a graph into a TVM Function

class nnvm.compiler.compile_engine.GraphKey(handle)

Key of a graph compilation context

nnvm.compiler.compile_engine.graph_key(graph, inputs, target)

Construct a new graph key.

Parameters:
  • graph (Graph) – The computation graph structure
  • inputs (list of Tensor(placeholder)) – The input requirement to the graph.
  • target (str) – The target of compilation.