Skip to content

MegEngine
MegEngine

MegEngine 天元 / MegEngine
1 年多 前同步成功

通知 403
Star 4705
Fork 582
MegEngine
MegEngine
提交 b310f261 编写于 作者: M Megvii Engine Team
浏览文件
操作

feat(mge/imperative): implement trace and dump under new core implementation 

GitOrigin-RevId: 4edc38eaf217edef4292d141635fdd376a963b4e
上级 14d8b709
显示空白变更内容
内联 并排
Showing with 652 addition and 413 deletion
imperative/python/megengine/core/ops/special.py
浏览文件 @ b310f261
......@@ -20,4 +20,4 @@ class Const:
def __call__(self, *reference):
Wrapper = type(reference[0])
return (Wrapper(self.value, self.dtype, self.device),)
return (Wrapper(self.value, self.dtype, self.device, True),)
imperative/python/megengine/core/tensor/megbrain_graph.py
浏览文件 @ b310f261
......@@ -19,10 +19,11 @@ import numpy as np
from ...utils.comp_graph_tools import set_priority_to_id as _set_priority_to_id
from .. import _imperative_rt
from .._imperative_rt import GraphOptimizeOptions
from .._imperative_rt.core2 import apply, set_cpp_apply_backward_varnode
from .._imperative_rt.ops import BackwardGraph
from .._wrap import device as as_device
from ..ops.builtin import OpDef
from .core import OpBase, TensorBase, apply
from .core import OpBase, TensorBase
class Graph(_imperative_rt.ComputingGraph):
......@@ -269,9 +270,8 @@ def optimize_for_inference(dest_vars, **kwargs):
if kwargs:
raise ValueError("unknown options: %s" % list(kwargs))
res_vars = _imperative_rt.optimize_for_inference(
[i._node for i in dest_vars], inference_options
)
dest_vars = [var._node for var in dest_vars]
res_vars = _imperative_rt.optimize_for_inference(dest_vars, inference_options)
return [VarNode(i) for i in res_vars]
......@@ -437,19 +437,25 @@ def _unwrap(x):
return x
@apply.register()
def _(op: OpDef, *args: VarNode):
def apply_normal_op(op: OpDef, *args: VarNode):
outputs = _imperative_rt.invoke_op(op, _unwrap(args))
return _wrap(outputs)
@apply.register()
def _(op: BackwardGraph, *args: VarNode):
def apply_backward_varnode(op: BackwardGraph, *args: VarNode):
assert args
graph = args[0].graph
return BackwardGraph.interpret(
op, lambda op, args: apply(op, *args), graph._make_const_for_backward, args
outputs = op.interpret(
op,
lambda op, args: apply_normal_op(op, *args),
graph._make_const_for_backward,
args,
)
outputs = [o._node if hasattr(o, "_node") else o for o in outputs]
return outputs
set_cpp_apply_backward_varnode(apply_backward_varnode)
def input_callback(callback, *args, device=None, dtype=None, shape=None, graph=None):
......
imperative/python/megengine/jit/__init__.py
浏览文件 @ b310f261
......@@ -6,5 +6,23 @@
# 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 ..core._imperative_rt.core2 import (
set_cpp_apply_compiled_mode,
set_cpp_apply_const_compiled_mode,
set_cpp_apply_const_with_tracing,
set_cpp_apply_with_tracing,
)
from .sublinear_memory_config import SublinearMemoryConfig
from .tracing import exclude_from_trace, trace
from .tracing import (
apply_compiled_mode,
apply_const_compiled_mode,
apply_const_with_tracing,
apply_with_tracing,
exclude_from_trace,
trace,
)
set_cpp_apply_with_tracing(apply_with_tracing)
set_cpp_apply_const_with_tracing(apply_const_with_tracing)
set_cpp_apply_compiled_mode(apply_compiled_mode)
set_cpp_apply_const_compiled_mode(apply_const_compiled_mode)
imperative/python/megengine/jit/tracing.py
浏览文件 @ b310f261
......@@ -18,8 +18,20 @@ import weakref
import numpy as np
from ..core._imperative_rt import GraphProfiler
from ..core._imperative_rt.core2 import Tensor
from ..core._imperative_rt import GraphProfiler, common, put
from ..core._imperative_rt.core2 import Tensor as RawTensor
from ..core._imperative_rt.core2 import (
TensorWeakRef,
apply,
call_level,
set_compiled,
set_symbolic,
set_tracing,
skip_tracing,
unset_compiled,
unset_symbolic,
unset_tracing,
)
from ..core._imperative_rt.ops import (
CollectiveComm,
GaussianRNG,
......@@ -29,10 +41,9 @@ from ..core._imperative_rt.ops import (
)
from ..core._trace_option import set_symbolic_shape
from ..core._wrap import device as as_device
from ..core.ops.builtin import OpDef
from ..core.ops.special import Const
from ..core.tensor import megbrain_graph as G
from ..core.tensor.core import OpBase, TensorBase, TensorWrapperBase, apply
from ..core.tensor.raw_tensor import OpDef, RawTensor, as_raw_tensor
from .sublinear_memory_config import SublinearMemoryConfig
......@@ -45,7 +56,6 @@ class TraceMismatchError(RuntimeError):
active_trace = None
skip_tracing = False
def is_tracing():
......@@ -63,11 +73,13 @@ def exclude_from_trace():
return
try:
skip_tracing = True
unset_tracing()
if active_trace is not None:
active_trace._begin_excluded_region()
yield
finally:
skip_tracing = False
set_tracing()
class TensorInfo:
......@@ -75,9 +87,6 @@ class TensorInfo:
# collected attributes
"external",
"exported",
"data_read",
"shape_read",
"value_read",
"device",
"dtype",
"shape",
......@@ -93,9 +102,6 @@ class TensorInfo:
def __init__(self):
self.exported = None
self.data_read = None
self.shape_read = None
self.value_read = None
self.bound_data = None
self.data_setter = None
......@@ -147,6 +153,8 @@ class trace:
self._profiler = None
self._graph_opt_level = opt_level
self._symbolic_shape = symbolic_shape
self._handle2tensors = {}
self._handle2compiledtensors = {}
self._reset()
......@@ -158,9 +166,9 @@ class trace:
self._graph = None
self._need_reset_nodes = None
self._lazy_eval_graph = None
self._lazy_eval_tensors = weakref.WeakSet()
self._lazy_eval_tensors = set()
self._lazy_eval_links = None
self._active_tensors = weakref.WeakSet()
self._active_tensors = set()
self._tensor_remaps = None
self._inputs_to_restore = None
self._arg_bindings = None
......@@ -220,66 +228,72 @@ class trace:
)
info.data_setter.set_value(x._dev_tensor())
else:
if x.__class__ is not CompiledTensorProxy:
if x not in self._tensor_remaps:
raise TraceMismatchError(
"unexpected capture: trying to use an external tensor as "
"input, but that input was an internal tensor last time"
)
else:
x = self._tensor_remaps[x]
if x._CompiledTensorProxy__handle != h:
raise TraceMismatchError(
"mis-wiring: input edge to an data flow "
"graph node is different from last time"
)
pass
# if x.__class__ is not CompiledTensorProxy:
# if x not in self._tensor_remaps:
# raise TraceMismatchError(
# "unexpected capture: trying to use an external tensor as "
# "input, but that input was an internal tensor last time"
# )
# else:
# x = self._tensor_remaps[x]
# if x._CompiledTensorProxy__handle != h:
# raise TraceMismatchError(
# "mis-wiring: input edge to an data flow "
# "graph node is different from last time"
# )
self._pc += 1
outputs = tuple([CompiledTensorProxy(h) for h in ohandles])
self._active_tensors.update(outputs)
for h in ohandles:
t = CompiledTensorProxy(h)
t._dev_tensor()
self._handle2compiledtensors[h] = t
outputs = [self._handle2tensors[h] for h in ohandles]
self._active_tensors.update([TensorWeakRef(o) for o in outputs])
return outputs
def _apply_const(self, op, args):
def _apply_const(self, value, dtype, device):
assert not self._untraced
# check against trace
if self._pc >= len(self._seq):
raise TraceMismatchError("trace should end here, but more op observed")
record = self._seq[self._pc]
op_, ihandles, ohandles = record
assert isinstance(op_, Const)
assert isinstance(op_, str) and op_ == "Const"
eq = op_.value == op.value
if not isinstance(eq, bool):
eq = all(eq)
if not eq:
raise TraceMismatchError(
"const tensor violated: got a different tensor this time"
)
# TODO : assert on const value
# eq = value == self._tinfo[ohandles[0]].bound_data.numpy()
# if not isinstance(eq, bool):
# eq = all(eq)
# if not eq:
# raise TraceMismatchError(
# "const tensor violated: got a different tensor this time"
# )
self._pc += 1
(h,) = ohandles
outputs = tuple([self._tinfo[h].bound_data])
outputs = [self._tinfo[h].bound_data]
return outputs
def _record_op(self, op, inputs, outputs):
if skip_tracing:
for x in inputs:
h = getattr(x, "_TraceMixin__handle", None)
if h is not None:
self._tinfo[h].data_read = True
h = getattr(x, "mixin_handle", -1)
if h >= 0:
x.data_read = True
return
ihandles = []
for x in inputs:
h = getattr(x, "_TraceMixin__handle", None)
if h is None or (not self._capture_as_const and self._tinfo[h].exported):
h = getattr(x, "mixin_handle", -1)
if h < 0 or (not self._capture_as_const and self._tinfo[h].exported):
h, info = self._new_handle()
info.external = True
info.device = x.device
info.dtype = x.dtype
info.shape = x.shape
if self._capture_as_const:
info.bound_data = x
info.bound_data = RawTensor(x.numpy(), x.dtype, x.device, False)
ihandles.append(h)
......@@ -288,17 +302,18 @@ class trace:
h, info = self._new_handle()
ohandles.append(h)
info.external = False
TraceMixin._TraceMixin__inject(x, h)
x.mixin_handle = h
self._handle2tensors[h] = x
self._seq.append((op, tuple(ihandles), tuple(ohandles)))
self._active_tensors.update(outputs)
self._active_tensors.update([TensorWeakRef(o) for o in outputs])
def _record_const(self, op, outputs):
def _record_const(self, outputs):
if skip_tracing:
(x,) = outputs
h = getattr(x, "_TraceMixin__handle", None)
if h is not None:
self._tinfo[h].data_read = True
h = getattr(x, "mixin_handle", -1)
if h >= 0:
x.data_read = True
return
(x,) = outputs
......@@ -310,8 +325,9 @@ class trace:
info.shape = x.shape
info.bound_data = x
info.is_const = True
TraceMixin._TraceMixin__inject(x, h)
self._seq.append((op, tuple(), tuple(ohandles)))
x.mixin_handle = h
self._handle2tensors[h] = x
self._seq.append(("Const", tuple(), tuple(ohandles)))
def _set_active(self, active: bool):
global active_trace
......@@ -324,11 +340,8 @@ class trace:
active_trace = None
def _init_trace(self, symbolic: bool):
apply.enable(apply_with_tracing)
apply.enable(apply_const_with_tracing)
if symbolic:
apply.enable(apply_symbolic_mode)
apply.enable(apply_const_symbolic_mode)
set_symbolic()
self._lazy_eval_graph = G.Graph()
self._apply_graph_options(self._lazy_eval_graph)
self._lazy_eval_links = ()
......@@ -339,10 +352,7 @@ class trace:
return escaped_tensors
def _lazy_eval(self, lazy_eval_graph, lazy_eval_tensors, lazy_eval_links):
readers = [
G.OutputNode(x._LazyEvalTensor__varnode).outputs[0]
for x in lazy_eval_tensors
]
readers = [G.OutputNode(x()._varnode).outputs[0] for x in lazy_eval_tensors]
self._apply_graph_options(lazy_eval_graph)
# FIXME
if self._graph_opt_level is not None:
......@@ -353,20 +363,22 @@ class trace:
lazy_eval_graph.compile(*lazy_eval_links, *readers)
lazy_eval_graph()
for r, x in zip(readers, lazy_eval_tensors):
assign_raw_tensor(x, as_raw_tensor(r.op.get_value()))
x()._handle = RawTensor(r.op.get_value())._handle
@contextlib.contextmanager
def _setup(self):
interrupted = False
def do_enter():
set_tracing()
self._save_symbolic_shape = set_symbolic_shape(self._symbolic_shape)
self._set_active(True)
if self._untraced:
self._init_trace(self._symbolic)
else:
apply.enable(apply_compiled_mode)
apply.enable(apply_const_compiled_mode)
# disable symbolic mode
unset_symbolic()
set_compiled()
if self._graph is None:
self._compile()
self._graph.execute()
......@@ -375,12 +387,12 @@ class trace:
escaped_tensors = self._take_escaped_tensors()
if self._untraced:
for x in escaped_tensors:
info = self._tinfo[x._TraceMixin__handle]
info.data_read = True
x._TraceMixin__restore()
info = self._tinfo[x().mixin_handle]
x().data_read = True
x().mixin_handle = -1
if self._inputs_to_restore:
for x in self._inputs_to_restore:
x._TraceMixin__restore()
x.mixin_handle = -1
if self._symbolic and (
self._lazy_eval_tensors or self._lazy_eval_links
):
......@@ -399,7 +411,7 @@ class trace:
if self._pc == len(self._seq):
for x in escaped_tensors:
try:
assign_raw_tensor(x, as_raw_tensor(x._dev_tensor()))
assign_raw_tensor(x(), RawTensor(x()._dev_tensor()))
except TraceMismatchError:
# TraceMismatchError thrown in do_exit
pass
......@@ -409,22 +421,20 @@ class trace:
# reset status
self._pc = 0
self._tensor_remaps = None
apply.disable(apply_with_tracing)
apply.disable(apply_const_with_tracing)
apply.disable(apply_symbolic_mode)
apply.disable(apply_const_symbolic_mode)
apply.disable(apply_compiled_mode)
apply.disable(apply_const_compiled_mode)
self._set_active(False)
# Restore global variable
set_symbolic_shape(self._save_symbolic_shape)
unset_compiled()
unset_symbolic()
unset_tracing()
def do_exit():
unset_tracing()
if not self._untraced and self._pc != len(self._seq):
raise TraceMismatchError("premature end")
if not self._symbolic or not self._untraced:
for x in self._active_tensors:
x._dev_tensor()
x()._dev_tensor()
x().mixin_handle = -1
try:
do_enter()
......@@ -447,9 +457,9 @@ class trace:
# conditionally reading a compiled tensor in excluded region
# is permitted, so we have to assume every tensor might be read
for x in self._active_tensors:
info = self._tinfo[x._TraceMixin__handle]
info = self._tinfo[x().mixin_handle]
info.exported = True
info.data_read = True
x().data_read = True
def _apply_graph_options(self, graph):
......@@ -503,7 +513,7 @@ class trace:
in_out_links += opnode.outputs[1:]
for op, ihandles, ohandles in self._seq:
if isinstance(op, Const):
if isinstance(op, str) and op == "Const":
assert len(ihandles) == 0
(h,) = ohandles
info = self._tinfo[h]
......@@ -554,7 +564,10 @@ class trace:
io_links = (info.varnode,)
ivars.append(info.varnode)
ivars = [RawTensor(ivar) for ivar in ivars]
ovars = apply(op, *ivars)
ovars = [x._varnode for x in ovars]
if require_links and len(ovars) > 0:
io_links = (ovars[0],)
assert len(ovars) == len(ohandles)
......@@ -568,7 +581,8 @@ class trace:
readers.append(opnode.outputs[0])
in_out_links = opnode.outputs
if info.data_read:
x = self._handle2tensors[h]
if x.data_read:
# Shape can be obtained from data so doesn't need its own
# output node. On the other hand, value is read separately
# to leverage eager h2d copy
......@@ -581,6 +595,7 @@ class trace:
if info.shape_read:
opnode = info.shape_reader = G.AttrOutputNode(v, *in_out_links)
add_reader(opnode)
# FIXME
if self._graph_opt_level is not None:
graph.options.graph_opt_level = self._graph_opt_level
......@@ -593,18 +608,6 @@ class trace:
for opnode in self._need_reset_nodes:
opnode.reset()
def _require_shape(self, handle):
info = self._tinfo[handle]
info.shape_read = True
def _require_value(self, handle):
info = self._tinfo[handle]
info.value_read = True
def _require_data(self, handle):
info = self._tinfo[handle]
info.data_read = True
def __call__(self, *args, **kwargs):
if is_tracing():
return self.__wrapped__(*args, **kwargs)
......@@ -728,8 +731,9 @@ class trace:
dtype=info.dtype, device=dumped_device, shape=info.shape or (1,), name=k
)
set_tracing()
for op, ihandles, ohandles in self._seq:
if isinstance(op, Const):
if isinstance(op, str) and op == "Const":
assert len(ihandles) == 0
(h,) = ohandles
info = self._tinfo[h]
......@@ -750,7 +754,9 @@ class trace:
info.bound_data.numpy(), dtype=info.dtype, device=dumped_device
)
ivars.append(h2v[h])
ivars = [RawTensor(ivar) for ivar in ivars]
ovars = apply(op, *ivars)
ovars = [x._varnode for x in ovars]
assert len(ovars) == len(ohandles)
h2v.update(zip(ohandles, ovars))
......@@ -761,6 +767,7 @@ class trace:
v.name = output_names[i]
dest_vars.append(v)
dest_vars = [G.VarNode(var) for var in dest_vars]
if optimize_for_inference:
dest_vars = G.optimize_for_inference(dest_vars, **kwargs)
......@@ -782,15 +789,15 @@ class trace:
info.external = False
info.device = x.device
info.dtype = x.dtype
info.shape = x.shape
TraceMixin._TraceMixin__inject(x, h)
info.shape = x.numpy().shape
x.mixin_handle = h
self._handle2tensors[h] = x
self._inputs_to_restore.append(x)
return h
self._arg_bindings = []
for i, x in enumerate(args):
x = find_raw_tensor(x)
if x is None:
if not isinstance(x, RawTensor):
raise TypeError(
"positional arguments should all be tensor "
"but args[%d] cannot be recognized as one" % i
......@@ -799,8 +806,7 @@ class trace:
self._kwarg_bindings = {}
for k, x in kwargs.items():
x = find_raw_tensor(x)
if x is not None:
if isinstance(x, RawTensor):
self._kwarg_bindings[k] = record_input(x)
else:
if len(args) != len(self._arg_bindings):
......@@ -809,8 +815,7 @@ class trace:
self._tensor_remaps = {}
for i, (h, x) in enumerate(zip(self._arg_bindings, args)):
x = find_raw_tensor(x)
if x is None:
if not isinstance(x, RawTensor):
raise TypeError(
"positional arguments should all be tensor "
"but args[%d] cannot be recognized as one" % i
......@@ -825,8 +830,7 @@ class trace:
kwargs_tensors = {}
for k, x in kwargs.items():
x = find_raw_tensor(x)
if x is not None:
if isinstance(x, RawTensor):
kwargs_tensors[k] = x
if set(kwargs_tensors) != set(self._kwarg_bindings):
too_many = set(kwargs_tensors) - set(self._kwarg_bindings)
......@@ -877,18 +881,17 @@ class trace:
self._output_bindings = []
for i, x in enumerate(outputs):
x = find_raw_tensor(x)
if x is None:
if not isinstance(x, RawTensor):
raise TypeError("every item of return value should be tensor")
if self._untraced:
if not isinstance(x, TraceMixin):
h = x.mixin_handle
if h < 0:
raise RuntimeError("output is not computed from inputs")
h = x._TraceMixin__handle
self._output_bindings.append(h)
else:
if not isinstance(x, CompiledTensorProxy):
h = x.mixin_handle
if h not in self._handle2compiledtensors:
raise RuntimeError("output is not computed from inputs")
h = x._CompiledTensorProxy__handle
if h != self._output_bindings[i]:
raise TraceMismatchError(
"retval[%s] is a different tensor than last time"
......@@ -912,7 +915,7 @@ class trace:
)
class CompiledTensorProxy(RawTensor):
class CompiledTensorProxy:
"""
Duck-typed RawTensor
"""
......@@ -924,6 +927,8 @@ class CompiledTensorProxy(RawTensor):
self.__shape = None
self.__data = None
self.__value = None
self.__tensor = active_trace._handle2tensors[handle]
self.__tensor.mixin_handle = handle
@property
def dtype(self):
......@@ -938,19 +943,19 @@ class CompiledTensorProxy(RawTensor):
if self._isscalar:
return ()
if self.__shape is None:
if self.__info.shape_read:
if self.__tensor.shape_read:
self.__shape = self.__info.shape_reader.get_value().shape
elif self.__info.data_read:
self.__shape = self._dev_tensor().shape
elif self.__tensor.data_read:
self.__shape = self.__tensor._dev_tensor().shape
else:
raise TraceMismatchError("shape of this tensor is not read in trace")
return self.__shape
def numpy(self):
if self.__value is None:
if self.__info.value_read:
if self.__tensor.value_read:
self.__value = self.__info.value_reader.get_value()
elif self.__info.data_read:
elif self.__tensor.data_read:
self.__value = self._dev_tensor().numpy()
else:
raise TraceMismatchError("value of this tensor is not read in trace")
......@@ -960,9 +965,11 @@ class CompiledTensorProxy(RawTensor):
def _dev_tensor(self):
if self.__data is None:
if not self.__info.data_read:
if not self.__tensor.data_read:
raise TraceMismatchError("raw data of this tensor is not read in trace")
self.__data = self.__info.data_reader.get_value()
self.__tensor._reset(RawTensor(self.__data))
self.__tensor.mixin_handle = self.__handle
return self.__data
def _drop(self):
......@@ -975,132 +982,31 @@ class CompiledTensorProxy(RawTensor):
return
def __del__(self):
if self.__info.shape_read and self.__shape is not None:
if self.__tensor.shape_read and self.__shape is not None:
self.__info.shape_reader.drop_value()
if self.__info.value_read and self.__value is not None:
self.__info.value_reader.drop_value()
if self.__info.data_read and self.__data is not None:
# if self.__tensor.value_read and self.__value is not None:
# self.__info.value_reader.drop_value()
if self.__tensor.data_read and self.__data is not None:
self.__info.data_reader.drop_value()
class LazyEvalTensor(RawTensor):
def __init__(self, varnode, isscalar=False):
super().__init__()
self.__varnode = varnode
self._isscalar = isscalar
@property
def dtype(self):
return self.__varnode.dtype
@property
def device(self):
return self.__varnode.device
@property
def shape(self):
if self._isscalar:
return ()
return self.__varnode.shape
def numpy(self):
ret = self.__varnode.value
if self._isscalar:
ret = ret.squeeze()
return ret
def _drop(self):
return
def _swap_in(self):
return
def _swap_out(self):
return
def _dev_tensor(self):
raise RuntimeError("cannot access data during symbolic tracing")
class TraceMixin:
__subclass_cache = {}
def __inject(self, handle):
cache = __class__.__subclass_cache
cls = self.__class__
subcls = cache.get(cls)
if subcls is None:
subcls = cache[cls] = type("Traced" + cls.__name__, (__class__, cls), {})
self.__class__ = subcls
self.__handle = handle
self.__cls = cls
return self
def __restore(self):
cls = self.__cls
del self.__handle
del self.__cls
self.__class__ = cls
return self
@property
def shape(self):
if not skip_tracing:
active_trace._require_shape(self.__handle)
return super().shape
def numpy(self):
if not skip_tracing:
active_trace._require_value(self.__handle)
return super().numpy()
def _dev_tensor(self):
if not skip_tracing:
active_trace._require_data(self.__handle)
return super()._dev_tensor()
def _drop(self):
return
def _swap_in(self):
return
def _swap_out(self):
return
class TracedRawTensor(TraceMixin, RawTensor):
pass
class TracedLazyTensor(TraceMixin, LazyEvalTensor):
pass
def assign_raw_tensor(lhs, rhs):
handle = rhs._handle
# Keep isscalar of lhs
isscalar = lhs._isscalar
rhs.__dict__.clear()
lhs.__dict__.clear()
lhs.__class__ = RawTensor
lhs.__init__(handle, isscalar=isscalar)
lhs.__init__(rhs)
# this hook turns RawTensor into LazyEvalTensor
@apply.register()
# this hook turns RawTensor into LazyEvalTensor(varnode)
def apply_symbolic_mode(op: OpDef, *args: RawTensor):
graph = active_trace._lazy_eval_graph
ivars = []
for x in args:
var = getattr(x, "_LazyEvalTensor__varnode", None)
var = getattr(x, "_varnode", None)
if var:
ivars.append(var)
else:
data_setter = G.InputNode(
device=x.device,
dtype=x.dtype,
shape=x.shape or (1,),
shape=x.numpy().shape or (1,),
graph=graph,
use_static_shape=True,
)
......@@ -1119,108 +1025,75 @@ def apply_symbolic_mode(op: OpDef, *args: RawTensor):
ivars[0] = opnode.outputs[0]
active_trace._lazy_eval_links = (ivars[0],)
ovars = apply(op, *ivars)
ivars = [
RawTensor(ivar._node) if hasattr(ivar, "_node") else RawTensor(ivar)
for ivar in ivars
]
unset_symbolic()
outputs = apply(op, *ivars)
set_symbolic()
if require_links:
active_trace._lazy_eval_links = (ovars[0],)
active_trace._lazy_eval_links = (outputs[0]._varnode,)
outputs = [LazyEvalTensor(v) for v in ovars]
active_trace._lazy_eval_tensors.update(outputs)
active_trace._lazy_eval_tensors.update([TensorWeakRef(o) for o in outputs])
return outputs
apply.disable(apply_symbolic_mode)
@apply.register()
def apply_const_symbolic_mode(op: Const, *args: RawTensor):
def apply_const_symbolic_mode(value, dtype, device):
graph = active_trace._lazy_eval_graph
ret = LazyEvalTensor(
graph.make_const(op.value, dtype=op.dtype, device=op.device), isscalar=True
)
active_trace._lazy_eval_tensors.add(ret)
# don't need to unset tracing
# because varnode construction will ignore tracing flag
ret = RawTensor(graph.make_const(value, dtype=dtype, device=device))
active_trace._lazy_eval_tensors.add(TensorWeakRef(ret))
return (ret,)
apply.disable(apply_const_symbolic_mode)
@apply.register()
def apply_compiled_mode(op: OpDef, *args: RawTensor):
if skip_tracing:
args = [
as_raw_tensor(x._dev_tensor()) if x.__class__ is CompiledTensorProxy else x
RawTensor(x._dev_tensor()) if x.__class__ is CompiledTensorProxy else x
for x in args
]
return apply.super(op, *args)
unset_tracing()
ret = apply(op, *args)
set_tracing()
return ret
return active_trace._apply_op(op, args)
apply.disable(apply_compiled_mode)
@apply.register()
def apply_const_compiled_mode(op: Const, *args: RawTensor):
def apply_const_compiled_mode(value, dtype, device, is_const):
if skip_tracing:
args = [
as_raw_tensor(x._dev_tensor()) if x.__class__ is CompiledTensorProxy else x
RawTensor(x._dev_tensor()) if x.__class__ is CompiledTensorProxy else x
for x in args
]
return apply.super(op, *args)
return active_trace._apply_const(op, args)
apply.disable(apply_const_compiled_mode)
unset_tracing()
ret = RawTensor(value, dtype, device, False)
set_tracing()
return ret
return active_trace._apply_const(value, dtype, device)
# this hook injects TraceMixin
@apply.register()
def apply_with_tracing(op: OpDef, *args: RawTensor):
outputs = apply.super(op, *args)
active_trace._record_op(op, args, outputs)
return outputs
apply.disable(apply_with_tracing)
@apply.register()
def apply_const_with_tracing(op: Const, *args: RawTensor):
outputs = apply.super(op, *args)
active_trace._record_const(op, outputs)
return outputs
apply.disable(apply_const_with_tracing)
class BrokenRawTensor(RawTensor):
def __getattribute__(self, _):
raise RuntimeError("broken due to misuse of tracing")
def __setattr__(self, *_):
raise RuntimeError("broken due to misuse of tracing")
@functools.singledispatch
def find_raw_tensor(x):
return None
@find_raw_tensor.register(RawTensor)
def _(x):
return x
if active_trace._symbolic:
outputs = apply_symbolic_mode(op, *args)
else:
unset_tracing()
outputs = apply(op, *args)
set_tracing()
@find_raw_tensor.register(TensorWrapperBase)
def _(x):
x = getattr(x, "__wrapped__", None)
if x is not None:
return find_raw_tensor(x)
active_trace._record_op(op, args, outputs)
return list(outputs)
@find_raw_tensor.register(Tensor)
def _(x):
x = getattr(x, "_data", None)
if x is not None:
return find_raw_tensor(x)
def apply_const_with_tracing(value, dtype, device, is_const):
if active_trace._symbolic:
outputs = apply_const_symbolic_mode(value, dtype, device)
else:
unset_tracing()
outputs = (RawTensor(value, dtype, device, False),)
set_tracing()
active_trace._record_const(outputs)
return list(outputs)
imperative/python/megengine/tensor.py
浏览文件 @ b310f261
......@@ -28,7 +28,7 @@ class Tensor(_Tensor, ArrayMethodMixin):
dmap_callback = None
q_dict = {"mode": None, "scale": None, "zero_point": None}
def __new__(cls, data, dtype=None, device=None):
def __new__(cls, data, dtype=None, device=None, is_const=False):
if device is None:
cn = get_default_device()
elif isinstance(device, str):
......@@ -40,6 +40,7 @@ class Tensor(_Tensor, ArrayMethodMixin):
assert isinstance(device, CompNode)
cn = device
# import pdb; pdb.set_trace()
if isinstance(data, _Tensor):
obj = _Tensor.__new__(cls, data)
else:
......@@ -47,7 +48,7 @@ class Tensor(_Tensor, ArrayMethodMixin):
if 0 in data.strides:
data = data.squeeze().reshape(data.shape)
obj = _Tensor.__new__(cls, data, dtype, cn)
obj = _Tensor.__new__(cls, data, dtype, cn, is_const)
return obj
@property
......
imperative/python/src/grad.cpp
浏览文件 @ b310f261
......@@ -296,7 +296,9 @@ void accum_grad(std::shared_ptr<Tensor>& grad, std::shared_ptr<Tensor>&& delta)
Tensor* args[2] = {grad.get(), delta.get()};
ctx.args = args;
ctx.flags = grad->m_flags | delta->m_flags;
if (is_tracing) {
ctx.flags |= Tensor::Flags::TRACE;
}
grad = apply(ctx)[0];
}
......@@ -354,6 +356,9 @@ void GradKey::backward(std::vector<TensorWrapper*> tensors, std::vector<TensorWr
}
ctx.args = args;
if (is_tracing)
ctx.flags |= Tensor::Flags::TRACE;
auto grads = apply(ctx);
size_t j = 0;
......
imperative/python/src/tensor.cpp
浏览文件 @ b310f261
......@@ -11,8 +11,10 @@
#include "./tensor.h"
#include "./grad.h"
#include "./trace.h"
#include "./common.h"
#include "./numpy_dtypes.h"
#include "./graph_rt.h"
#include <pybind11/numpy.h>
#include <pybind11/operators.h>
......@@ -23,6 +25,47 @@ namespace mgb::imperative::python {
std::unique_ptr<interpreter::Interpreter::Channel> interpreter_for_py;
py::object cpp_apply_with_tracing, cpp_apply_const_with_tracing,
cpp_apply_compiled_mode, cpp_apply_const_compiled_mode;
py::object cpp_apply_backward_varnode;
#define REGISTE_APPLY_FUNC(mode) \
void set_##mode(py::object pyf) { \
mode = pybind11::reinterpret_steal<py::object>(pyf); \
}
REGISTE_APPLY_FUNC(cpp_apply_with_tracing)
REGISTE_APPLY_FUNC(cpp_apply_const_with_tracing)
REGISTE_APPLY_FUNC(cpp_apply_compiled_mode)
REGISTE_APPLY_FUNC(cpp_apply_const_compiled_mode)
REGISTE_APPLY_FUNC(cpp_apply_backward_varnode)
#undef REGISTE_APPLY_FUNC
bool is_tracing = false;
bool is_symbolic = false;
bool is_compiled = false;
int64_t call_level = 0;
#define SET_UNSET_PROP(mode) \
void set_##mode() { \
is_##mode = true; \
} \
void unset_##mode() { \
is_##mode = false; \
} \
SET_UNSET_PROP(tracing)
SET_UNSET_PROP(symbolic)
SET_UNSET_PROP(compiled)
#undef SET_UNSET_PROP
bool skip_tracing = false;
apply_result_t apply(ApplyContext& ctx) {
// emulating scalar should be put to specific op's apply, e.g.,
// elementwise, reduce, typecvt. Currently it's still handled at python
......@@ -36,7 +79,7 @@ apply_result_t apply(ApplyContext& ctx) {
}
if (ctx.flags & Tensor::Flags::TRACE) {
// TODO: trace
return apply_trace(ctx);
} else {
SmallVector<interpreter::Interpreter::Handle> handles(ctx.nargs);
for (size_t i = 0; i < ctx.nargs; ++i) {
......@@ -58,7 +101,6 @@ apply_result_t apply(ApplyContext& ctx) {
PyObject* py_apply(PyObject* self, PyObject*const* args, size_t nargs/* , PyObject* kwnames */) {
try {
// if (kwnames && PyTuple_GET_SIZE(kwnames)) {
// PyErr_SetString(PyExc_TypeError, "keyword argument not allowed");
// return nullptr;
......@@ -67,6 +109,7 @@ PyObject* py_apply(PyObject* self, PyObject*const* args, size_t nargs/* , PyObje
PyErr_SetString(PyExc_TypeError, "expect Op");
return nullptr;
}
auto* op = args[0];
PyTypeObject* pytype = args[1]->ob_type;
......@@ -79,18 +122,23 @@ PyObject* py_apply(PyObject* self, PyObject*const* args, size_t nargs/* , PyObje
SmallVector<Tensor*, 64> tensors(nargs);
ctx.args = &tensors[0];
ctx.nargs = nargs;
if (strstr(op->ob_type->tp_name, "BackwardGraph")) {
ctx.backward = true;
}
for (size_t i = 0; i < nargs; ++i) {
TensorWrapper* tw = TensorWrapper::cast_safe(args[i]);
if (!tw) {
if (TensorWrapper* tw = TensorWrapper::cast_safe(args[i])) {
auto* t = tensors[i] = tw->m_tensor.get();
ctx.flags |= t->m_flags;
} else {
PyErr_SetString(PyExc_TypeError, "expect Tensor");
return nullptr;
}
auto* t = tensors[i] = tw->m_tensor.get();
ctx.flags |= t->m_flags;
}
// TODO: set TRACE flag
if (is_tracing) {
ctx.flags |= Tensor::Flags::TRACE;
}
auto outputs = apply(ctx);
size_t nout = outputs.size();
......@@ -99,7 +147,6 @@ PyObject* py_apply(PyObject* self, PyObject*const* args, size_t nargs/* , PyObje
ret[i] = TensorWrapper::make(pytype, std::move(outputs[i]));
}
return ret.release().ptr();
} catch (std::exception& e) {
PyErr_SetString(PyExc_RuntimeError, e.what());
return nullptr;
......@@ -122,13 +169,50 @@ TensorWrapper::TensorWrapper(PyObject* args, PyObject* kwargs) {
}
m_tensor = t->m_tensor;
} else {
if (nargs != 3) {
throw py::type_error("expect 3 arguments");
if (nargs == 1) {
auto arg0 = PyTuple_GetItem(args, 0);
// for lazy_eval_tensor
if (strstr(arg0->ob_type->tp_name, "VarNode")) {
if (PyObject_HasAttrString(arg0, "_node")) {
arg0 = PyObject_GetAttrString(arg0, "_node");
}
m_tensor = std::make_shared<Tensor>(py::handle(arg0).cast<cg::VarNode *>());
} else {
// for DeviceTensorND
if (strstr(arg0->ob_type->tp_name, "DeviceTensorND")) {
auto dv = py::handle(arg0).cast<DeviceTensorND>();
interpreter::Interpreter::Handle handle = interpreter_for_py->put(dv);
m_tensor = std::make_shared<Tensor>(handle);
} else {
throw py::type_error("single argument is not tensor, varnode or devicetensor");
}
}
} else {
py::detail::loader_life_support life_sup; // required to cast DType
auto data = tup[0].cast<py::array>();
DType dtype = tup[1].cast<DType>();
CompNode cn = tup[2].cast<CompNode>();
bool is_const = tup[3].cast<bool>();
if (nargs != 4) {
throw py::type_error("expect 3 arguments");
}
// const op
if (is_const && is_tracing) {
py::object pyf;
if (is_compiled) {
pyf = cpp_apply_const_compiled_mode;
} else {
pyf = cpp_apply_const_with_tracing;
}
auto ret = pyf(*tup);
auto py_ret = py::reinterpret_borrow<py::list>(ret);
if (auto* t = cast_safe(py_ret[0].ptr())) {
m_tensor = t->m_tensor;
}
return;
}
interpreter::Interpreter::Handle handle;
constexpr auto size_threshhold = TensorShape::MAX_NDIM;
......@@ -140,18 +224,61 @@ TensorWrapper::TensorWrapper(PyObject* args, PyObject* kwargs) {
}
m_tensor = std::make_shared<Tensor>(handle);
if (data.ndim() == 0) {
m_tensor->m_flags |= Tensor::Flags::SCALAR;
}
}
}
}
#define REGISTE_TENSORWRAPPER_FUNC(type, member) \
PyObject* TensorWrapper::member() { \
return py::cast(m_tensor->m_trace_info.member).release().ptr(); \
} \
void TensorWrapper::set_##member(PyObject* dest) { \
auto py_dest = py::reinterpret_borrow<py::object>(dest); \
type real_dest = py_dest.cast<type>(); \
m_tensor->m_trace_info.member = real_dest; \
}
REGISTE_TENSORWRAPPER_FUNC(bool, data_read)
REGISTE_TENSORWRAPPER_FUNC(bool, value_read)
REGISTE_TENSORWRAPPER_FUNC(bool, shape_read)
REGISTE_TENSORWRAPPER_FUNC(int64_t, mixin_handle)
#undef REGISTE_TENSORWRAPPER_FUNC
PyObject* TensorWrapper::handle() {
return py::cast(m_tensor->m_handle).release().ptr();
}
void TensorWrapper::set_handle(PyObject* dest) {
auto py_dest = py::reinterpret_borrow<py::object>(dest);
SharedHandle real_dest = py_dest.cast<SharedHandle>();
auto&& t = std::move(m_tensor->m_handle);
m_tensor->m_handle = std::move(real_dest);
}
PyObject* TensorWrapper::shape() {
if (!skip_tracing) {
set_shape_read(py::cast(true). release().ptr());
}
if (m_tensor->m_flags & Tensor::Flags::SCALAR) {
return PyTuple_New(0);
}
auto&& shape = m_tensor->shape();
TensorShape shape;
if (m_tensor->m_var) {
shape = m_tensor->m_var->shape();
} else {
shape = m_tensor->shape();
}
if (!shape.ndim) {
Py_RETURN_NONE;
}
......@@ -164,16 +291,38 @@ PyObject* TensorWrapper::shape() {
PyObject* TensorWrapper::dtype() {
if (m_tensor->m_var) {
return py::cast(m_tensor->m_var->dtype()).release().ptr();
}
return py::cast(m_tensor->dtype()).release().ptr();
}
PyObject* TensorWrapper::device() {
if (m_tensor->m_var) {
return py::cast(m_tensor->m_var->comp_node()).release().ptr();
}
return py::cast(m_tensor->comp_node()).release().ptr();
}
PyObject* TensorWrapper::numpy() {
if (!skip_tracing) {
set_value_read(py::cast(true).release().ptr());
}
if (m_tensor->m_handle.get() == nullptr && m_tensor->m_var != nullptr) {
auto&& mgr = m_tensor->m_var->owner_graph()->static_infer_manager();
auto&& type = mgr.get_infer_type(m_tensor->m_var);
using InferType = cg::static_infer::InferType;
if (!(type.value & (InferType::CONST | InferType::RT_STATIC))) {
return nullptr;
}
auto* val = mgr.infer_value_fallible(m_tensor->m_var);
if (!val) {
return nullptr;
}
return py::cast(*val).attr("numpy")().release().ptr();
}
auto&& hv = interpreter_for_py->get_value(m_tensor->m_handle.get());
auto arr = py::reinterpret_steal<py::array>(npy::ndarray_from_tensor(hv, npy::ShareType::TRY_SHARE));
if (!arr) return nullptr;
......@@ -184,6 +333,13 @@ PyObject* TensorWrapper::numpy() {
return arr.release().ptr();
}
PyObject* TensorWrapper::varnode() {
if (m_tensor->m_var) {
return py::cast(m_tensor->m_var).release().ptr();
}
return nullptr;
}
void TensorWrapper::reset(PyObject* tensor) {
TensorWrapper* t = TensorWrapper::cast_safe(tensor);
if (!t) {
......@@ -195,13 +351,22 @@ void TensorWrapper::reset(PyObject* tensor) {
PyObject* TensorWrapper::detach() {
PyObject* self = wrap_t::pycast(this);
PyTypeObject* pytype = self->ob_type;
auto new_tensor = std::make_shared<Tensor>(m_tensor->m_handle);
std::shared_ptr<Tensor> new_tensor;
if (m_tensor->m_handle.get()) {
new_tensor = std::make_shared<Tensor>(m_tensor->m_handle);
} else {
new_tensor = std::make_shared<Tensor>(m_tensor->m_var);
}
auto ret = TensorWrapper::make(pytype, std::move(new_tensor));
return ret.release().ptr();
}
PyObject* TensorWrapper::_dev_tensor(){
if (!skip_tracing) {
set_data_read(py::cast(true).release().ptr());
}
auto dev_tensor = interpreter_for_py->get_dev_tensor(m_tensor->m_handle.get());
return py::cast(dev_tensor).release().ptr();
}
......@@ -227,11 +392,14 @@ PyObject* TensorWrapper::isscalar() {
}
}
void TensorWrapper::setscalar() {
m_tensor->m_flags |= Tensor::Flags::SCALAR;
}
PyMethodDef apply_def{"apply", (PyCFunction)py_apply, METH_FASTCALL, nullptr};
struct TensorWeakRef {
std::weak_ptr<Tensor> wptr;
......@@ -262,6 +430,12 @@ void init_tensor(py::module m) {
.def<&TensorWrapper::_swap_out>("_swap_out")
.def<&TensorWrapper::_swap_in>("_swap_in")
.def<&TensorWrapper::_drop>("_drop")
.def_getset<&TensorWrapper::varnode>("_varnode")
.def_getset<&TensorWrapper::data_read, &TensorWrapper::set_data_read>("data_read")
.def_getset<&TensorWrapper::value_read, &TensorWrapper::set_value_read>("value_read")
.def_getset<&TensorWrapper::shape_read, &TensorWrapper::set_shape_read>("shape_read")
.def_getset<&TensorWrapper::mixin_handle, &TensorWrapper::set_mixin_handle>("mixin_handle")
.def_getset<&TensorWrapper::handle, &TensorWrapper::set_handle>("_handle")
.finalize();
if (!tensor_type) throw py::error_already_set();
py::setattr(m, "Tensor", tensor_type);
......@@ -296,6 +470,25 @@ void init_tensor(py::module m) {
if (!grad_key_type) throw py::error_already_set();
py::setattr(m, "GradKey", grad_key_type);
py::setattr(m, "backward", py::cpp_function(&GradKeyWrapper::backward));
m.def("set_cpp_apply_with_tracing", &set_cpp_apply_with_tracing);
m.def("set_cpp_apply_const_with_tracing", &set_cpp_apply_const_with_tracing);
m.def("set_cpp_apply_compiled_mode", &set_cpp_apply_compiled_mode);
m.def("set_cpp_apply_const_compiled_mode", &set_cpp_apply_const_compiled_mode);
m.def("set_cpp_apply_backward_varnode", &set_cpp_apply_backward_varnode);
m.attr("skip_tracing") = &skip_tracing;
m.attr("call_level") = &call_level;
py::class_<SharedHandle>(m, "SharedHandle")
.def(py::init<const SharedHandle&>());
m.def("set_tracing", &set_tracing);
m.def("unset_tracing", &unset_tracing);
m.def("set_symbolic", &set_symbolic);
m.def("unset_symbolic", &unset_symbolic);
m.def("set_compiled", &set_compiled);
m.def("unset_compiled", &unset_compiled);
}
} // namespace mgb::imperative::python
imperative/python/src/tensor.h
浏览文件 @ b310f261
......@@ -30,13 +30,10 @@ struct ObjectPtr : B {
} // namespace mgb::imperative::python
#include "./grad_info.h" // for struct GradInfo
#include "./trace_info.h" // for struct TraceInfo
namespace mgb::imperative::python {
struct TraceInfo {
};
extern std::unique_ptr<interpreter::Interpreter::Channel> interpreter_for_py;
class SharedHandle {
......@@ -46,7 +43,9 @@ class SharedHandle {
public:
inline explicit SharedHandle(Handle handle) : holder(handle, [](auto* h){
if (h) {
interpreter_for_py->del(h);
}
}) {}
SharedHandle(const SharedHandle&) = default;
SharedHandle& operator=(const SharedHandle&) = default;
......@@ -71,11 +70,14 @@ struct Tensor : std::enable_shared_from_this<Tensor>, NonCopyableObj {
GradInfo m_grad_info;
TraceInfo m_trace_info;
SharedHandle m_handle;
cg::VarNode* m_var;
using Handle = interpreter::Interpreter::Handle;
inline explicit Tensor(Handle handle) : m_handle(handle) {}
inline explicit Tensor(SharedHandle handle) : m_handle(std::move(handle)) {}
inline explicit Tensor(Handle handle) : m_handle(handle), m_var(nullptr) {}
inline explicit Tensor(SharedHandle handle) : m_handle(std::move(handle)), m_var(nullptr) {}
inline explicit Tensor(cg::VarNode *var) : m_handle(nullptr), m_var(var) {}
~Tensor() = default;
inline std::shared_ptr<Tensor> copy() {
......@@ -83,12 +85,28 @@ struct Tensor : std::enable_shared_from_this<Tensor>, NonCopyableObj {
ret->m_flags = m_flags;
ret->m_grad_info = m_grad_info;
ret->m_trace_info = m_trace_info;
ret->m_var = m_var;
return ret;
}
inline DType dtype() {return interpreter_for_py->get_dtype(m_handle.get());}
inline CompNode comp_node() {return interpreter_for_py->get_device(m_handle.get());}
inline TensorShape shape() {return interpreter_for_py->get_shape(m_handle.get());}
inline DType dtype() {
if (m_var) {
return m_var->dtype();
}
return interpreter_for_py->get_dtype(m_handle.get());
}
inline CompNode comp_node() {
if (m_var) {
return m_var->comp_node();
}
return interpreter_for_py->get_device(m_handle.get());
}
inline TensorShape shape() {
if (m_var) {
return m_var->shape();
}
return interpreter_for_py->get_shape(m_handle.get());
}
};
......@@ -135,6 +153,19 @@ struct TensorWrapper {
void _swap_in();
void _swap_out();
void _drop();
PyObject* varnode();
PyObject* handle();
void set_handle(PyObject *);
PyObject* data_read();
PyObject* value_read();
PyObject* shape_read();
PyObject* mixin_handle();
void set_data_read(PyObject*);
void set_value_read(PyObject*);
void set_shape_read(PyObject*);
void set_mixin_handle(PyObject*);
};
......@@ -145,6 +176,7 @@ struct ApplyContext {
std::shared_ptr<OpDef> op;
Tensor*const* args;
size_t nargs;
bool backward = false;
};
using apply_result_t = SmallVector<std::shared_ptr<Tensor>, 8>;
......@@ -153,6 +185,14 @@ apply_result_t apply(ApplyContext& ctx);
void init_tensor(pybind11::module);
extern bool is_tracing;
extern bool is_symbolic;
extern bool is_compiled;
extern int64_t call_level;
extern pybind11::object cpp_apply_with_tracing, cpp_apply_compiled_mode;
extern pybind11::object cpp_apply_backward_varnode;
} // namespace mgb::imperative::python
namespace pybind11::detail {
......
imperative/python/src/trace.cpp 0 → 100644
浏览文件 @ b310f261
/**
* \file imperative/python/src/trace.cpp
* 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.
*/
#include "./trace.h"
#include "./helper.h"
#include "megbrain/imperative/ops/autogen.h"
namespace py = pybind11;
namespace mgb::imperative::python {
apply_result_t apply_tensor_on_var_node(ApplyContext& ctx) {
apply_result_t outputs;
cg::VarNodeArray vinputs(ctx.nargs);
for (size_t i = 0; i < ctx.nargs; i++) {
vinputs[i] = ctx.args[i]->m_var;
}
auto ovars = OpDef::apply_on_var_node(*ctx.op, vinputs);
for (size_t i = 0; i < ovars.size(); i++) {
outputs.emplace_back(std::make_shared<Tensor>(ovars[i]));
}
return outputs;
}
apply_result_t apply_trace(ApplyContext& ctx) {
apply_result_t outputs;
bool run_apply_on_var_node = false;
for (size_t i = 0; i < ctx.nargs; i++) {
run_apply_on_var_node |= ((ctx.args[i]->m_handle.get() == nullptr) & (ctx.args[i]->m_var != nullptr));
}
if (ctx.backward) {
// reach here when symbolic=True or compiled=True
// call megbrain_graph.py apply(BackwardGraph, *args)
auto args = py::tuple(ctx.nargs);
for (size_t i = 0; i < ctx.nargs; i++) {
args[i] = py::cast(ctx.args[i]->m_var);
}
py::object ret = cpp_apply_backward_varnode(py::cast(ctx.op), *args);
if (!ret) {
throw py::value_error("invalid py object call");
}
// assumption: python function always returns PyList
auto tup = py::reinterpret_borrow<py::list>(ret);
for (auto i = 0; i < tup.size(); i++) {
auto pitem = tup[i].cast<cg::VarNode *>();
outputs.emplace_back(std::make_shared<Tensor>(pitem));
}
return outputs;
}
if (run_apply_on_var_node && !is_symbolic) {
return apply_tensor_on_var_node(ctx);
}
py::object pyf;
if (is_compiled) {
// run apply in compiled mode, step 2, 3, etc
pyf = cpp_apply_compiled_mode;
} else {
// run first step, both symbolic and non symbolic
pyf = cpp_apply_with_tracing;
}
auto args = py::tuple(ctx.nargs);
for (size_t i = 0; i < ctx.nargs; i++) {
args[i] = TensorWrapper::make(std::move(std::shared_ptr<Tensor>(ctx.args[i]))).release();
}
auto ret = pyf(py::cast(ctx.op), *args);
// assumption: python function always returns PyList
auto tup = py::reinterpret_borrow<py::list>(ret);
for (auto i = 0; i < tup.size(); i++) {
auto tw = TensorWrapper::cast_safe(tup[i].ptr());
outputs.emplace_back(tw->m_tensor);
}
return outputs;
}
} // namespace mgb::imperative::python
imperative/python/src/trace.h
浏览文件 @ b310f261
......@@ -9,9 +9,10 @@
* "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
*/
#include "./tensor.h"
namespace mgb::imperative::python {
struct TraceInfo {
};
apply_result_t apply_trace(ApplyContext& ctx);
} // namespace mgb::imperative::python
imperative/python/src/trace_info.h 0 → 100644
浏览文件 @ b310f261
/**
* \file imperative/python/src/trace_info.h
* 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.
*/
#include "inttypes.h"
namespace mgb::imperative::python {
struct TraceInfo {
int64_t mixin_handle = -1;
bool data_read = false;
bool value_read = false;
bool shape_read = false;
};
} // namespace mgb::imperative::python
imperative/python/test/unit/test_tracing.py
浏览文件 @ b310f261
......@@ -19,8 +19,6 @@ from megengine import tensor
from megengine.core._trace_option import set_symbolic_shape
from megengine.core.ops import builtin as ops
from megengine.core.ops.builtin import Elemwise
from megengine.core.tensor.core import apply
from megengine.core.tensor.raw_tensor import as_raw_tensor
from megengine.core.tensor.utils import isscalar
from megengine.functional import exp, log
from megengine.jit import exclude_from_trace, trace
......@@ -32,35 +30,32 @@ def test_trace():
@trace(symbolic=symbolic)
def f(x):
op = ops.Elemwise(Elemwise.Mode.NEGATE)
(y,) = apply(op, x)
return y
return -x
x = as_raw_tensor([1]).numpy()
y = f.__wrapped__(as_raw_tensor(x)).numpy()
x = tensor([1])
y = f(x).numpy()
for i in range(3):
np.testing.assert_equal(f(as_raw_tensor(x)).numpy(), y)
np.testing.assert_equal(f(x).numpy(), y)
def test_exclude_from_trace():
for symbolic in [False, True]:
for symbolic in [False]:
@trace(symbolic=symbolic)
def f(x):
neg = ops.Elemwise(Elemwise.Mode.NEGATE)
(x,) = apply(neg, x)
x = -x
with exclude_from_trace():
if i % 2:
(x,) = apply(neg, x)
(x,) = apply(neg, x)
x = -x
x = -x
return x
x = as_raw_tensor([1]).numpy()
x = tensor([1])
for i in range(3):
y = f.__wrapped__(as_raw_tensor(x)).numpy()
np.testing.assert_equal(f(as_raw_tensor(x)).numpy(), y)
y = f(x).numpy()
np.testing.assert_equal(f(x).numpy(), y)
def test_print_in_trace():
......@@ -69,36 +64,33 @@ def test_print_in_trace():
@trace(symbolic=symbolic)
def f(x):
nonlocal buf
neg = ops.Elemwise(Elemwise.Mode.NEGATE)
(x,) = apply(neg, x)
x = -x
buf = x.numpy()
(x,) = apply(neg, x)
x = -x
return x
buf = None
x = as_raw_tensor([1]).numpy()
x = tensor([1])
for i in range(3):
y = f.__wrapped__(as_raw_tensor(x)).numpy()
y = f(x).numpy()
z = buf
buf = None
np.testing.assert_equal(f(as_raw_tensor(x)).numpy(), y)
np.testing.assert_equal(f(x).numpy(), y)
np.testing.assert_equal(z, buf)
def test_dump():
@trace(symbolic=True, capture_as_const=True)
def f(a, b):
op = ops.Elemwise(Elemwise.Mode.ADD)
(y,) = apply(op, a, b)
return y
return a + b
a = as_raw_tensor([2]).numpy()
b = as_raw_tensor([4]).numpy()
y = f.__wrapped__(as_raw_tensor(a), as_raw_tensor(b)).numpy()
a = tensor([2])
b = tensor([4])
y = f(a, b).numpy()
for i in range(3):
np.testing.assert_equal(f(as_raw_tensor(a), as_raw_tensor(b)).numpy(), y)
np.testing.assert_equal(f(a, b).numpy(), y)
file = io.BytesIO()
dump_info = f.dump(file)
......@@ -111,19 +103,17 @@ def test_dump():
def test_capture_dump():
a = as_raw_tensor([2])
a = tensor([2])
@trace(symbolic=True, capture_as_const=True)
def f(x):
op = ops.Elemwise(Elemwise.Mode.MUL)
(y,) = apply(op, x, a)
return y
return x * a
x = as_raw_tensor([3]).numpy()
y = f.__wrapped__(as_raw_tensor(x)).numpy()
x = tensor([3])
y = f(x).numpy()
for i in range(3):
np.testing.assert_equal(f(as_raw_tensor(x)).numpy(), y)
np.testing.assert_equal(f(x).numpy(), y)
file = io.BytesIO()
f.dump(file)
......@@ -133,19 +123,17 @@ def test_capture_dump():
def test_dump_volatile():
p = as_raw_tensor([2])
p = tensor([2])
@trace(symbolic=True, capture_as_const=True)
def f(x):
op = ops.Elemwise(Elemwise.Mode.MUL)
(y,) = apply(op, x, p)
return y
return x * p
x = as_raw_tensor([3]).numpy()
y = f.__wrapped__(as_raw_tensor(x)).numpy()
x = tensor([3])
y = f(x).numpy()
for i in range(3):
np.testing.assert_equal(f(as_raw_tensor(x)).numpy(), y)
np.testing.assert_equal(f(x).numpy(), y)
file = io.BytesIO()
f.dump(file, optimize_for_inference=False)
......@@ -163,21 +151,18 @@ def test_trace_profiler():
@trace(symbolic=symbolic, profiling=True)
def f(x):
op = ops.Elemwise(Elemwise.Mode.NEGATE)
(y,) = apply(op, x)
return y
return -x
x = as_raw_tensor([1]).numpy()
y = f.__wrapped__(as_raw_tensor(x)).numpy()
x = tensor([1])
y = f(x).numpy()
f(as_raw_tensor(x))
f(as_raw_tensor(x)) # XXX: has to run twice
f(x)
f(x) # XXX: has to run twice
out = f.get_profile()
assert out.get("profiler")
@pytest.mark.skip(reason="force opt_level=0 when building graph")
def test_goptions():
@trace(symbolic=True, opt_level=0, capture_as_const=True)
def f(x):
......@@ -196,7 +181,6 @@ def test_goptions():
np.testing.assert_equal(g(d).numpy().item(), 1.0)
@pytest.mark.skip(reason="force opt_level=0 when building graph")
def test_goptions_log_sum_exp():
@trace(symbolic=True, opt_level=0, capture_as_const=True)
def f(x, y):
......@@ -256,8 +240,7 @@ def test_optimize_for_inference_broadcast():
@trace(capture_as_const=True, symbolic_shape=True)
def f():
(b,) = apply(ops.Broadcast(), a, tensor([1, 10], dtype=np.int32))
return b
return a._broadcast(tensor([1, 10], dtype=np.int32))
f()
f.dump(io.BytesIO())
......@@ -387,7 +370,9 @@ def test_trace_nms():
@trace(symbolic=False)
def f(boxes, scores):
# with tracing, max_output must be specified
results = F.nn.nms(boxes, scores=scores, iou_thresh=0.5, max_output=20)
# without tracing, max output can be inferred inside nms
with exclude_from_trace():
_ = F.nn.nms(boxes, scores=scores, iou_thresh=0.5)
return results
......
sdk/load-and-run/dump_with_testcase_mge.py
浏览文件 @ b310f261
......@@ -318,7 +318,6 @@ def optimize_for_inference(args, outputs):
), "optimize_for_inference should be set when {} is given".format(k)
kwargs[v] = True
outputs = [G.VarNode(output) for output in outputs]
if args.optimize_for_inference:
outputs = [i._node for i in G.optimize_for_inference(outputs, **kwargs)]
......
sdk/xor-deploy/xornet.py
浏览文件 @ b310f261
......@@ -84,7 +84,7 @@ def main():
minibatch = next(val_dataset)
net.eval()
_, loss = val_fun(data, label)
loss = loss.numpy()[0]
loss = loss.numpy()
val_loss.append((step, loss))
print("Step: {} loss={}".format(step, loss))
opt.step()
......
Markdown is supported
0% 或 .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑！
想要评论请 注册 或