diff --git a/imperative/python/megengine/__init__.py b/imperative/python/megengine/__init__.py index fa1d6962355c322baf1e6925134cb586c165fa4d..083032fd4acf9703f66e92b4233a6a1703a96a55 100644 --- a/imperative/python/megengine/__init__.py +++ b/imperative/python/megengine/__init__.py @@ -76,7 +76,7 @@ from .logger import enable_debug_log, get_logger, set_log_file, set_log_level from .serialization import load, save from .tensor import Parameter, Tensor, tensor from .version import __version__ -from .core import cgtools +from .utils import comp_graph_tools as cgtools _set_fork_exec_path_for_timed_func( sys.executable, diff --git a/imperative/python/megengine/core/__init__.py b/imperative/python/megengine/core/__init__.py index 50d29e9ebfa39fa1ee10b3169d0b00f1a12ebeb4..4fd130bc002d8574c1e367e2361370db56a65132 100644 --- a/imperative/python/megengine/core/__init__.py +++ b/imperative/python/megengine/core/__init__.py @@ -11,4 +11,3 @@ import sys from .tensor import Tensor from .tensor.megbrain_graph import Graph -from .utils import comp_graph_tools as cgtools diff --git a/imperative/python/megengine/core/tensor/megbrain_graph.py b/imperative/python/megengine/core/tensor/megbrain_graph.py index df20a6cd745b20786930acb6be4a3d58fafb2721..9b4707b74b67d7d9d385d66dcfad74816e45348a 100644 --- a/imperative/python/megengine/core/tensor/megbrain_graph.py +++ b/imperative/python/megengine/core/tensor/megbrain_graph.py @@ -358,7 +358,7 @@ CompGraphLoadResult = collections.namedtuple( def load_graph(fpath): """Load a serialized computing graph from file. - :parma fpath: Path or Handle for the output file + :param fpath: Path or Handle of the input file :return: An instance of namedtuple :class:`CompGraphLoadResult`, whose fields are: diff --git a/imperative/python/megengine/core/utils/__init__.py b/imperative/python/megengine/core/utils/__init__.py deleted file mode 100644 index 200f5ab2fd8f948d207499d7603660e8b3a2924c..0000000000000000000000000000000000000000 --- a/imperative/python/megengine/core/utils/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -# -*- coding: utf-8 -*- -# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") -# -# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -from .comp_graph_tools import * diff --git a/imperative/python/megengine/core/utils/comp_graph_tools.py b/imperative/python/megengine/utils/comp_graph_tools.py similarity index 86% rename from imperative/python/megengine/core/utils/comp_graph_tools.py rename to imperative/python/megengine/utils/comp_graph_tools.py index ceffcc7efeda7afb860308e4e390f2e4bfe298d9..13726c86dbf154eab6c8ffd9b7e84492ed24b8a0 100644 --- a/imperative/python/megengine/core/utils/comp_graph_tools.py +++ b/imperative/python/megengine/utils/comp_graph_tools.py @@ -8,8 +8,12 @@ import collections from typing import Dict, List -from .. import _imperative_rt -from .._imperative_rt import OperatorNode, VarNode +import numpy + +from ..core import _imperative_rt +from ..core._imperative_rt import OperatorNode, VarNode +from ..core.tensor import megbrain_graph as G +from ..core.tensor.raw_tensor import as_raw_tensor def get_dep_vars(var: VarNode, var_type: str = None) -> List[VarNode]: @@ -251,3 +255,33 @@ def set_priority_to_id(dest_vars): assert isinstance(i, VarNode) dest_vec.append(i) _imperative_rt.graph._set_priority_to_id(dest_vec) + + +def load_and_inference(file, inp_data_list: List[numpy.ndarray]) -> List[numpy.ndarray]: + """Load a serialized computing graph and run inference with input data. + + :param file: Path or Handle of the input file. + :param inp_data_list: List of input data. + :return: List of inference results. + + """ + *_, out_list = G.load_graph(file) + inputs = get_dep_vars(out_list, "Host2DeviceCopy") + replace_dict = {} + inp_node_list = [] + for i in inputs: + inp_node = G.InputNode( + device="xpux", dtype=inputs[0].dtype, graph=inputs[0].graph + ) + replace_dict[i] = inp_node.outputs[0] + inp_node_list.append(inp_node) + new_out = replace_vars(out_list, replace_dict) + out_node_list = [G.OutputNode(i) for i in new_out] + new_out_list = [i.outputs[0] for i in out_node_list] + cg = new_out_list[0].graph + func = cg.compile(new_out_list) + for node, value in zip(inp_node_list, inp_data_list): + node.set_value(as_raw_tensor(value)._dev_tensor()) + func.execute() + out_data_list = [o.get_value().numpy() for o in out_node_list] + return out_data_list diff --git a/imperative/python/test/unit/test_tracing.py b/imperative/python/test/unit/test_tracing.py index 5c63f8b3c3153e007196330ced018ba6d091fdaf..3a6a28d5fcc110877a5f51348e61541445292c2d 100644 --- a/imperative/python/test/unit/test_tracing.py +++ b/imperative/python/test/unit/test_tracing.py @@ -22,29 +22,6 @@ from megengine.functional import exp, log from megengine.jit import exclude_from_trace, trace -def load_and_inference(file, inp_data): - cg, _, out_list = G.load_graph(file) - inputs = cgtools.get_dep_vars(out_list, "Host2DeviceCopy") - replace_dict = {} - inp_node_list = [] - for i in inputs: - inp_node = G.InputNode( - device="xpux", dtype=inputs[0].dtype, graph=inputs[0].graph - ) - replace_dict[i] = inp_node.outputs[0] - inp_node_list.append(inp_node) - new_out = cgtools.replace_vars(out_list, replace_dict) - out_node_list = [G.OutputNode(i) for i in new_out] - new_out_list = [i.outputs[0] for i in out_node_list] - new_cg = new_out_list[0].graph - func = new_cg.compile(new_out_list) - for node, value in zip(inp_node_list, inp_data): - node.set_value(as_raw_tensor(value)._dev_tensor()) - func.execute() - out_data_list = [o.get_value().numpy() for o in out_node_list] - return out_data_list - - def test_trace(): for symbolic in [False, True]: @@ -124,7 +101,7 @@ def test_dump(): 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]) + result = cgtools.load_and_inference(file, [a, b]) np.testing.assert_equal(result[0], y) @@ -146,7 +123,7 @@ def test_capture_dump(): file = io.BytesIO() f.dump(file) file.seek(0) - result = load_and_inference(file, [x]) + result = cgtools.load_and_inference(file, [x]) np.testing.assert_equal(result[0], y) diff --git a/sdk/load-and-run/dump_with_testcase_imperative.py b/sdk/load-and-run/dump_with_testcase_imperative.py index 8ebd12747547feb882631cd8db142a72e3691150..b7a347ad690ecc7171e6a0290885ae471ee458f3 100755 --- a/sdk/load-and-run/dump_with_testcase_imperative.py +++ b/sdk/load-and-run/dump_with_testcase_imperative.py @@ -17,6 +17,7 @@ import numpy as np import megengine as mge import megengine.core._imperative_rt as rt import megengine.core.tensor.megbrain_graph as G +from megengine.core.tensor.megbrain_graph import VarNode from megengine import cgtools from megengine.core.ops import builtin from megengine.core.tensor.core import apply @@ -488,7 +489,8 @@ def main(): with open(args.output, "wb") as fout: fout.write(b"mgbtest0") fout.write(struct.pack("I", len(feeds["testcases"]))) - fout.write(rt.dump_graph(output_mgbvars)) + dump_content, _ = G.dump_graph([VarNode(i) for i in output_mgbvars]) + fout.write(dump_content) def make_dev_tensor(value, dtype=None, device=None): return as_raw_tensor(value, dtype=dtype, device=device)._dev_tensor() @@ -507,7 +509,8 @@ def main(): testcase.keys() ) with open(args.output, "ab") as fout: - fout.write(G.dump_graph(*output_mgbvars)) + dump_content, _ = G.dump_graph(output_mgbvars) + fout.write(dump_content)