Skip to content

Commit

Permalink
[DOCS] Add save_param_dict, readme (apache#42)
Browse files Browse the repository at this point in the history
  • Loading branch information
tqchen authored and sergei-mironov committed Aug 8, 2018
1 parent bca8909 commit c89811d
Show file tree
Hide file tree
Showing 12 changed files with 412 additions and 49 deletions.
46 changes: 42 additions & 4 deletions nnvm/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -3,16 +3,54 @@
[![Build Status](https://travis-ci.org/dmlc/nnvm.svg?branch=master)](https://travis-ci.org/dmlc/nnvm)
[![GitHub license](http://dmlc.github.io/img/apache2.svg)](./LICENSE)

NNVM is a reusable computational graph optimization and compilation stack for deep learning systems.
NNVM provides modules to:
NNVM is a reusable computational graph optimization and compilation stack for deep learning systems. It provides modules to:

- Represent deep learning workloads from front-end frameworks via a graph IR.
- Optimize computation graphs to improve performance.
- Compile into executable modules and deploy to different hardware backends with minimum dependency.

NNVM is designed to add new frontend, operators and graph optimizations in a decentralized fashion without changing the core interface. NNVM is part of [TVM stack](https://github.com/dmlc/tvm), which provides an end to end IR compilation stack for deploying deep learning workloads into different hardware backends
NNVM is designed to add new frontend, operators and graph optimizations in a decentralized fashion without changing the core interface. NNVM is part of [TVM stack](https://github.com/dmlc/tvm). NNVM compiler toolchain can target hardware backends supported by TVM.
The compiled module can be deployed to server, mobile, embedded devices and browsers with minimum dependency, in languages including c++, python, javascript, java, objective-c.

The following code snippet demonstrates the general workflow of nnvm compiler toolchain.

```python
import tvm
from tvm.contrib import graph_runtime, rpc
import nnvm.frontend
import nnvm.compiler

# get model from frameworks
# change xyz to supported framework name.
graph, params = nnvm.frontend.from_xyz(...)

# optimize and compile the graph to get a deployable module
# target can be "opencl", "llvm", "metal" or any target supported by tvm
target = "cuda"
graph, lib, params = nnvm.compiler.build(
graph, target, shape={"data", data_shape}, params=params)

# deploy and run on gpu(0)
module = graph_runtime.create(graph, lib, tvm.gpu(0))
module.set_input(**params)
output = tvm.nd.empty(out_shape, ctx=tvm.gpu(0))
for data_array in dataset:
module.set_input("data", data_array)
module.run()
module.get_output(0, output)

# deploy to remote mobile/rasp/browser with minimum tvm rpc runtime
# useful for quick experiments on mobile devices
remote = rpc.connect(remote_host, remote_port)
lib.export_library("mylib.so")
remote.upload("mylib.so")
rlib = rpc.load_module("mylib.so")
# run on remote device
rmodule = graph_runtime.create(graph, rlib, remote.gpu(0))
rmodule.set_input(**params)
rmodule.run()
```

## Links
- [TinyFlow](https://github.com/tqchen/tinyflow) on how you can use NNVM to build a TensorFlow like API.
- [Apache MXNet](http://mxnet.io/) uses NNVM as a backend.

4 changes: 4 additions & 0 deletions nnvm/docs/api/python/compiler.rst
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,10 @@ nnvm.compiler

.. autofunction:: nnvm.compiler.build_config

.. autofunction:: nnvm.compiler.save_param_dict

.. autofunction:: nnvm.compiler.load_param_dict

.. autofunction:: nnvm.compiler.optimize

.. automodule:: nnvm.compiler.graph_util
Expand Down
4 changes: 3 additions & 1 deletion nnvm/python/nnvm/compiler/__init__.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
"""NNVM compiler toolchain.
User only need to use :any:`build` and :any:`build_config` to do the compilation.
User only need to use :any:`build` and :any:`build_config` to do the compilation,
and :any:`save_param_dict` to save the parameters into bytes.
The other APIs are for more advanced interaction with the compiler toolchain.
"""
from __future__ import absolute_import
Expand All @@ -10,6 +11,7 @@
from . import build_module
from . build_module import build, optimize, build_config
from . compile_engine import engine, graph_key
from . param_dict import save_param_dict, load_param_dict

from .. import symbol as _symbol
from .. import graph as _graph
Expand Down
32 changes: 28 additions & 4 deletions nnvm/python/nnvm/compiler/build_module.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
from .. import graph as _graph

OPT_PASS_LEVEL = {
"SimplifyInference": 2,
"SimplifyInference": 0,
"PrecomputePrune": 2,
"OpFusion": 1
}
Expand All @@ -26,6 +26,7 @@ class BuildConfig(object):
current = None
defaults = {
"opt_level": 2,
"add_pass": None,
}
def __init__(self, **kwargs):
self._old_scope = None
Expand Down Expand Up @@ -53,6 +54,23 @@ def __exit__(self, ptype, value, trace):
assert self._old_scope
BuildConfig.current = self._old_scope

def pass_enabled(self, pass_name):
"""Get whether pass is enabled.
Parameters
----------
pass_name : str
The optimization pass name
Returns
-------
enabled : bool
Whether pass is enabled.
"""
if self.add_pass and pass_name in self.add_pass:
return True
return self.opt_level >= OPT_PASS_LEVEL[pass_name]


BuildConfig.current = BuildConfig()

Expand All @@ -64,6 +82,9 @@ def build_config(**kwargs):
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: BuildConfig
Expand Down Expand Up @@ -120,7 +141,7 @@ def optimize(graph, shape, dtype="float32"):
"""
# pylint: disable=unused-argument
cfg = BuildConfig.current
if cfg.opt_level >= OPT_PASS_LEVEL["SimplifyInference"]:
if cfg.pass_enabled("SimplifyInference"):
graph = graph_attr.set_shape_inputs(graph, shape)
graph = graph.apply(["InferShape", "SimplifyInference"])
return graph
Expand Down Expand Up @@ -182,14 +203,17 @@ def build(graph, target, shape, dtype="float32", params=None):
# Apply optimization
graph = optimize(graph, shape, dtype)
# Precompute prune
if params and cfg.opt_level >= OPT_PASS_LEVEL["PrecomputePrune"]:
if params and cfg.pass_enabled("PrecomputePrune"):
graph, params = precompute_prune(graph, params)
shape, dtype = _update_shape_dtype(shape, dtype, params)
# Operator Fusion and generatiom
graph = graph_attr.set_shape_inputs(graph, shape)
graph = graph_attr.set_dtype_inputs(graph, dtype)
graph._set_json_attr("target", target, "str")
graph._set_json_attr("opt_level", cfg.opt_level, "int")
if cfg.pass_enabled("OpFusion"):
graph._set_json_attr("opt_level", 1, "int")
else:
graph._set_json_attr("opt_level", 0, "int")
graph = graph.apply("InferShape").apply("InferType")
graph = graph.apply("GraphFusePartition").apply("GraphFuseCompile")
libmod = graph_attr._move_out_module(graph, "module")
Expand Down
66 changes: 66 additions & 0 deletions nnvm/python/nnvm/compiler/param_dict.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
"""Helper utility to save parameter dict"""
import tvm

_save_param_dict = tvm.get_global_func("nnvm.compiler._save_param_dict")
_load_param_dict = tvm.get_global_func("nnvm.compiler._load_param_dict")

def 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: bytearray
Serialized parameters.
Examples
--------
.. code-block:: python
# 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)
"""
args = []
for k, v in params.items():
args.append(k)
args.append(tvm.nd.array(v))
return _save_param_dict(*args)


def load_param_dict(param_bytes):
"""Load parameter dictionary to binary bytes.
Parameters
----------
param_bytes: bytearray
Serialized parameters.
Returns
-------
params : dict of str to NDArray
The parameter dictionary.
"""
if isinstance(param_bytes, (bytes, str)):
param_bytes = bytearray(param_bytes)
load_mod = _load_param_dict(param_bytes)
size = load_mod(0)
param_dict = {}
for i in range(size):
key = load_mod(1, i)
dltensor_handle = load_mod(2, i)
param_dict[key] = tvm.nd.NDArray(dltensor_handle, False)
return param_dict
32 changes: 30 additions & 2 deletions nnvm/python/nnvm/graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
from ._base import c_array, c_str, nn_uint, py_str, string_types
from ._base import GraphHandle, SymbolHandle
from ._base import check_call
from .symbol import Symbol, Group as _Group
from .symbol import Variable, Symbol, Group as _Group

class GraphIndex(object):
"""Index for quickly accessing graph attributes.
Expand Down Expand Up @@ -174,9 +174,19 @@ def symbol(self):
check_call(_LIB.NNGraphGetSymbol(self.handle, ctypes.byref(shandle)))
return Symbol(shandle)

def json(self):
"""Get JSON representation of the graph
Returns
-------
json : str
JSON representation of the graph
"""
return self.apply("SaveJSON").json_attr("json")

def _tvm_graph_json(self):
"""Get TVM graph json"""
return self.apply("SaveJSON").json_attr("json")
return self.json()

@property
def index(self):
Expand Down Expand Up @@ -225,6 +235,24 @@ def apply(self, passes):
return Graph(ghandle)


def load_json(json_str):
"""Create a new graph by loading from json
Parameters
----------
json_str : str
The json string
Returns
-------
graph : Graph
The loaded graph
"""
ret = create(Variable("x"))
ret._set_json_attr("json", json_str)
return ret.apply("LoadJSON")


def create(symbol):
"""Create a new graph from symbol.
Expand Down
38 changes: 1 addition & 37 deletions nnvm/src/compiler/graph_fuse.cc
Original file line number Diff line number Diff line change
Expand Up @@ -15,46 +15,10 @@
#include <tvm/lowered_func.h>
#include <dmlc/parameter.h>
#include "./compile_engine.h"
#include "../../tvm/src/runtime/graph/graph_runtime.h"
#include "./graph_runtime.h"

namespace nnvm {
namespace compiler {


struct TVMOpParam : public dmlc::Parameter<TVMOpParam> {
std::string func_name;
uint32_t num_inputs;
uint32_t num_outputs;
uint32_t flatten_data;

DMLC_DECLARE_PARAMETER(TVMOpParam) {
DMLC_DECLARE_FIELD(func_name);
DMLC_DECLARE_FIELD(num_inputs).set_default(1);
DMLC_DECLARE_FIELD(num_outputs).set_default(1);
DMLC_DECLARE_FIELD(flatten_data).set_default(0);
}
};

DMLC_REGISTER_PARAMETER(TVMOpParam);

// parser
inline void TVMOpParamParser(nnvm::NodeAttrs* attrs) {
TVMOpParam param;
param.Init(attrs->dict);
attrs->parsed = std::move(param);
}

NNVM_REGISTER_OP(tvm_op)
.set_attr_parser(TVMOpParamParser)
.set_num_inputs([](const NodeAttrs& attrs) {
const TVMOpParam& param = nnvm::get<TVMOpParam>(attrs.parsed);
return param.num_inputs;
})
.set_num_outputs([](const NodeAttrs& attrs) {
const TVMOpParam& param = nnvm::get<TVMOpParam>(attrs.parsed);
return param.num_outputs;
});

using namespace tvm;

// The single fuse rule.
Expand Down
Loading

0 comments on commit c89811d

Please sign in to comment.