From fc74d5f2c3c50ad779c81f4a5d799ab841ce7bd2 Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Fri, 11 Sep 2020 16:32:57 +0800 Subject: [PATCH] feat(mge/dump): add more options for megbrain_graph.dump_graph GitOrigin-RevId: 848e3c87c183e716c750a2744f0dc27a50e5186a --- .../megengine/core/tensor/megbrain_graph.py | 279 ++++++++++++------ imperative/python/megengine/jit/tracing.py | 4 +- imperative/python/src/graph_rt.cpp | 111 +++++-- imperative/python/test/unit/test_tracing.py | 9 +- 4 files changed, 279 insertions(+), 124 deletions(-) diff --git a/imperative/python/megengine/core/tensor/megbrain_graph.py b/imperative/python/megengine/core/tensor/megbrain_graph.py index 965fe9c69..df20a6cd7 100644 --- a/imperative/python/megengine/core/tensor/megbrain_graph.py +++ b/imperative/python/megengine/core/tensor/megbrain_graph.py @@ -11,6 +11,7 @@ import json import threading import weakref from concurrent.futures import Future, ThreadPoolExecutor +from typing import Dict, List, Union import numpy as np @@ -85,6 +86,97 @@ class Graph(_imperative_rt.ComputingGraph): return self._wrap(_imperative_rt.make_h2d(self, device, dtype, shape, name)) +class VarNode(TensorBase): + def __init__(self, node: _imperative_rt.VarNode): + self._node = node + if hasattr(self.graph, "_var_cache"): + self.graph._var_cache[node] = self + + @property + def graph(self) -> Graph: + return self._node.graph + + @property + def op(self): + if hasattr(self.graph, "_wrap"): + return self.graph._wrap(self._node.owner) + else: + return self._node.owner + + @property + def name(self): + return self._node.name + + @property + def id(self): + return self._node.id + + @name.setter + def name(self, name): + self._node.name = name + + @property + def dtype(self): + return self._node.dtype + + @property + def device(self): + return as_device(self._node.comp_node) + + @property + def shape(self): + return self._node.shape + + @property + def value(self): + return self._node.value + + +class OpNode: + def __init__(self, node: _imperative_rt.OperatorNode): + self._node = node + if hasattr(self.graph, "_op_cache"): + self.graph._op_cache[node] = self + + @property + def graph(self) -> Graph: + return self._node.graph + + @property + def name(self): + return self._node.name + + @property + def id(self): + return self._node.id + + @name.setter + def name(self, name): + self._node.name = name + + @property + def inputs(self): + if hasattr(self.graph, "_wrap"): + return tuple(map(self.graph._wrap, self._node.inputs)) + else: + return self._node.inputs + + @property + def outputs(self): + if hasattr(self.graph, "_wrap"): + return tuple(map(self.graph._wrap, self._node.outputs)) + else: + return self._node.outputs + + @property + def params(self): + return json.loads(self._node.params) + + @property + def type(self): + return self._node.type + + def optimize_for_inference(dest_vars, **kwargs): r"""Applies optimize_for_inference pass for computing graph. @@ -162,8 +254,100 @@ def optimize_for_inference(dest_vars, **kwargs): return [VarNode(i) for i in res_vars] -def dump_graph(*args): - return _imperative_rt.dump_graph([i._node for i in args]) +CompGraphDumpResult = collections.namedtuple( + "CompGraphDumpResult", + [ + "nr_opr", + "tot_bytes", + "tensor_value_bytes", + "content_hash", + "inputs", + "outputs", + "params", + ], +) + + +def dump_graph( + output_vars: Union[Dict[str, VarNode], List[VarNode]], + *, + keep_var_name: int = 1, + keep_param_name: bool = False, + keep_opr_priority: bool = False, + strip_info_file=None, +): + """serialize the computing graph of `output_vars` and get byte result. + + :param output_vars: output variables which are the graph's end point. + + .. note:: + + The underlying C++ API only accepts a var list. If a dict is given, + the vars would be renamed to the given names. + + :param keep_var_name: level for keeping variable names: + + * 0: none of the names are kept + * 1: (default)keep names of output vars + * 2: keep names of all (output and internal) vars + :param keep_param_name: whether to keep param names, so param values can be + easily manipulated after loading model + :param keep_opr_priority: whether to keep priority setting for operators + :param strip_info_file: a string for path or a file handler. if is not None, + then the dump information for code strip would be written to ``strip_info_file`` + :return: dump result as byte string, and an instance of namedtuple + :class:`CompGraphDumpResult`, whose fields are: + + * ``nr_opr`` number of operators dumped + * ``tot_bytes`` total bytes for the whole graph + * ``tensor_value_bytes`` bytes consumed for dumping tensor values + * ``inputs`` names of input tensors + * ``params`` list of names of dumped params + * ``outputs`` names of output vars + """ + ov = [] + if isinstance(output_vars, dict): + used_vars = set() + for name, var in output_vars.items(): + assert isinstance(var, VarNode), "bad output var: {!r}".format(var) + assert var.id not in used_vars, ( + "var name is associated with a var object, so we can not have " + "two names given to the same var: {}".format(var) + ) + used_vars.add(var.id) + var.name = name + ov.append(var._node) + else: + for var in output_vars: + assert isinstance(var, VarNode), "bad output var: {!r}".format(var) + ov.append(var._node) + + stat = [] + inputs = [] + outputs = [] + params = [] + + dump_content = _imperative_rt.dump_graph( + ov, + keep_var_name, + keep_param_name, + keep_opr_priority, + stat, + inputs, + outputs, + params, + ) + + dump_info = CompGraphDumpResult(*stat, inputs, outputs, params) + + if strip_info_file is not None: + if isinstance(strip_info_file, str): + strip_info_file = open(strip_info_file, "w") + strip_info = json.loads(_imperative_rt.get_info_for_strip(ov)) + strip_info["hash"] = dump_info.content_hash + json.dump(strip_info, strip_info_file) + + return dump_content, dump_info CompGraphLoadResult = collections.namedtuple( @@ -193,97 +377,6 @@ def load_graph(fpath): return CompGraphLoadResult(cg, dict(output_vars_map), output_vars_list) -class VarNode(TensorBase): - def __init__(self, node: _imperative_rt.VarNode): - self._node = node - if hasattr(self.graph, "_var_cache"): - self.graph._var_cache[node] = self - - @property - def graph(self) -> Graph: - return self._node.graph - - @property - def op(self): - if hasattr(self.graph, "_wrap"): - return self.graph._wrap(self._node.owner) - else: - return self._node.owner - - @property - def name(self): - return self._node.name - - @property - def id(self): - return self._node.id - - @name.setter - def name(self, name): - self._node.name = name - - @property - def dtype(self): - return self._node.dtype - - @property - def device(self): - return as_device(self._node.comp_node) - - @property - def shape(self): - return self._node.shape - - @property - def value(self): - return self._node.value - - -class OpNode: - def __init__(self, node: _imperative_rt.OperatorNode): - self._node = node - if hasattr(self.graph, "_op_cache"): - self.graph._op_cache[node] = self - - @property - def graph(self) -> Graph: - return self._node.graph - - @property - def name(self): - return self._node.name - - @property - def id(self): - return self._node.id - - @name.setter - def name(self, name): - self._node.name = name - - @property - def inputs(self): - if hasattr(self.graph, "_wrap"): - return tuple(map(self.graph._wrap, self._node.inputs)) - else: - return self._node.inputs - - @property - def outputs(self): - if hasattr(self.graph, "_wrap"): - return tuple(map(self.graph._wrap, self._node.outputs)) - else: - return self._node.outputs - - @property - def params(self): - return json.loads(self._node.params) - - @property - def type(self): - return self._node.type - - def _wrap(x): if isinstance(x, collections.abc.Sequence): return type(x)(map(_wrap, x)) diff --git a/imperative/python/megengine/jit/tracing.py b/imperative/python/megengine/jit/tracing.py index c1662fbf9..e219f607b 100644 --- a/imperative/python/megengine/jit/tracing.py +++ b/imperative/python/megengine/jit/tracing.py @@ -589,7 +589,9 @@ class trace: if isinstance(file, str): permission = "wb" if append == False else "ab" file = open(file, permission) - file.write(G.dump_graph(*dest_vars)) + dump_content, dump_info = G.dump_graph(dest_vars) + file.write(dump_content) + return dump_info def _process_inputs(self, *args, **kwargs): if self._untraced: diff --git a/imperative/python/src/graph_rt.cpp b/imperative/python/src/graph_rt.cpp index 98a1a49bc..4bc4c194b 100644 --- a/imperative/python/src/graph_rt.cpp +++ b/imperative/python/src/graph_rt.cpp @@ -27,6 +27,7 @@ namespace py = pybind11; using namespace mgb; using namespace imperative; +namespace ser = mgb::serialization; using _OptimizeForInferenceOptions = mgb::gopt::OptimizeForInferenceOptions; using _LayoutTransform = _OptimizeForInferenceOptions::LayoutTransform; @@ -183,7 +184,6 @@ void init_graph_rt(py::module m) { return "Opr:" + opr->name(); }); - py::class_(m, "AsyncExecutable") .def("execute", &cg::AsyncExecutable::execute, py::call_guard()) .def("wait", &cg::AsyncExecutable::wait, py::call_guard()); @@ -206,15 +206,6 @@ void init_graph_rt(py::module m) { })) .def("get", [](_CompGraphProfilerImpl& profiler) { return profiler._get_result(); }); - m.def("dump_graph", [](const std::vector& dest_vars) { - using namespace mgb::serialization; - std::vector buf; - auto dumper = GraphDumper::make(OutputFile::make_vector_proxy(&buf)); - SymbolVarArray symvars(dest_vars.begin(), dest_vars.end()); - dumper->dump(symvars); - return py::bytes(reinterpret_cast(&buf[0]), buf.size()); - }); - auto GraphOptimizeOptions = py::class_<_OptimizeForInferenceOptions>(m, "GraphOptimizeOptions") .def(py::init()) .def_readwrite("f16_io_f32_comp", &_OptimizeForInferenceOptions::f16_io_f32_comp) @@ -245,23 +236,91 @@ void init_graph_rt(py::module m) { return vars; }); + m.def("get_info_for_strip", [](const std::vector& dest_vars) { + std::unordered_set opr_types, dtype_names, elemwise_modes; + auto on_opr = [&](cg::OperatorNodeBase *opr) { + if (ser::GraphDumper::should_remove_in_dump(opr)) + return; + opr_types.insert(opr->dyn_typeinfo()->name); + for (auto i : opr->output()) + dtype_names.insert(i->dtype().name()); + if (opr->same_type()) { + auto mode = opr->cast_final().param().mode; + elemwise_modes.insert( + megdnn::Elemwise::ModeTrait::from_mode(mode).name); + } + }; + cg::DepOprIter opr_iter{on_opr}; + for (auto i : dest_vars) + opr_iter.add(i->owner_opr()); + + auto to_json = [](const std::unordered_set &v) { + std::vector vs(v.begin(), v.end()); + std::sort(vs.begin(), vs.end()); + auto ret = json::Array::make(); + for (auto &&i : vs) + ret->add(json::String::make(i)); + return ret; + }; + + return json::Object::make({ + {"opr_types", to_json(opr_types)}, + {"dtypes", to_json(dtype_names)}, + {"elemwise_modes", to_json(elemwise_modes)}, + }); + }); + + m.def("dump_graph", []( + const std::vector& dest_vars, + int keep_var_name, + bool keep_param_name, + bool keep_opr_priority, + py::list& stat, + py::list& inputs, + py::list& outputs, + py::list& params + ) { + std::vector buf; + auto dumper = ser::GraphDumper::make(ser::OutputFile::make_vector_proxy(&buf)); + SymbolVarArray symvars(dest_vars.begin(), dest_vars.end()); + + ser::GraphDumper::DumpConfig config{keep_var_name, keep_param_name, + keep_opr_priority}; + + auto rst = dumper->dump(symvars, config); + for (auto i : rst.inputs) { + inputs.append(py::cast(i)); + } + for (auto i : rst.outputs) { + outputs.append(py::cast(i)); + } + for (auto i : rst.params) { + params.append(py::cast(i)); + } + auto rst_stat = std::vector{ + rst.nr_opr, rst.tot_bytes, rst.tensor_value_bytes, rst.content_hash + }; + for (auto i : rst_stat) { + stat.append(py::cast(i)); + } + return py::bytes(reinterpret_cast(&buf[0]), buf.size()); + }); - m.def("load_graph", [](std::string& buf, py::list& _output_var_map, py::list& _output_var_list) { - using namespace mgb::serialization; - auto file = InputFile::make_mem_proxy(buf.c_str(), buf.length()); - auto format = GraphLoader::identify_graph_dump_format(*file); - auto loader = GraphLoader::make(std::move(file), format.val()); - GraphLoader::LoadConfig config; + m.def("load_graph", []( + std::string& buf, + py::list& output_var_map, + py::list& output_var_list + ) { + auto file = ser::InputFile::make_mem_proxy(buf.c_str(), buf.length()); + auto format = ser::GraphLoader::identify_graph_dump_format(*file); + auto loader = ser::GraphLoader::make(std::move(file), format.val()); + ser::GraphLoader::LoadConfig config; auto rst = loader->load(config); - std::vector> output_var_map; - SymbolVarArray output_var_list; - output_var_map = {rst.output_var_map.begin(), rst.output_var_map.end()}; - output_var_list = std::move(rst.output_var_list); - for (auto i : output_var_list){ - _output_var_list.append(i.node()); + for (auto i : rst.output_var_map) { + output_var_map.append(py::make_tuple(i.first, i.second.node())); } - for (auto i : output_var_map){ - _output_var_map.append(py::make_tuple(i.first,i.second.node())); + for (auto i : rst.output_var_list) { + output_var_list.append(i.node()); } std::unordered_map tensor2name; for (const auto& pair : rst.tensor_map) { @@ -277,8 +336,8 @@ void init_graph_rt(py::module m) { h2d.output(0)->name(*it->second); }; cg::DepOprIter iter{cb}; - for (const auto& var : output_var_list) { - iter.add(var.node()->owner_opr()); + for (const auto& var : rst.output_var_list) { + iter.add(var); } return rst.graph; diff --git a/imperative/python/test/unit/test_tracing.py b/imperative/python/test/unit/test_tracing.py index 78bf6c7c4..5c63f8b3c 100644 --- a/imperative/python/test/unit/test_tracing.py +++ b/imperative/python/test/unit/test_tracing.py @@ -12,12 +12,10 @@ from tempfile import mkstemp import numpy as np import pytest -import megengine -import megengine.module as M +import megengine.core.tensor.megbrain_graph as G from megengine import cgtools, tensor from megengine.core._trace_option import set_tensor_shape from megengine.core.ops import builtin as ops -from megengine.core.tensor import megbrain_graph as G from megengine.core.tensor.core import apply from megengine.core.tensor.raw_tensor import as_raw_tensor from megengine.functional import exp, log @@ -121,7 +119,10 @@ def test_dump(): np.testing.assert_equal(f(as_raw_tensor(a), as_raw_tensor(b)).numpy(), y) file = io.BytesIO() - f.dump(file) + dump_info = f.dump(file) + assert dump_info.nr_opr == 3 + np.testing.assert_equal(dump_info.inputs, ["h2d[0]", "h2d[2]"]) + np.testing.assert_equal(dump_info.outputs, ["ADD(h2d[0],h2d[2])[4]"]) file.seek(0) result = load_and_inference(file, [a, b]) np.testing.assert_equal(result[0], y) -- GitLab