From aed681d3256243f1547e0a6e10d4ee2ecfd983c0 Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Wed, 31 Mar 2021 18:04:30 +0800 Subject: [PATCH] feat(imperative/utils): optimize the naming rules GitOrigin-RevId: 329bac640aa6e2e3c981aa294a361684b982892e --- imperative/python/megengine/jit/tracing.py | 14 ++-- imperative/python/megengine/module/module.py | 8 +- imperative/python/megengine/tensor.py | 4 +- imperative/python/megengine/utils/naming.py | 74 +++++++++++-------- .../python/test/unit/jit/test_tracing.py | 4 +- .../test/unit/utils/test_dump_naming.py | 18 ++++- src/opr/impl/io.cpp | 2 +- 7 files changed, 74 insertions(+), 50 deletions(-) diff --git a/imperative/python/megengine/jit/tracing.py b/imperative/python/megengine/jit/tracing.py index 8fddd0420..669a659f0 100644 --- a/imperative/python/megengine/jit/tracing.py +++ b/imperative/python/megengine/jit/tracing.py @@ -40,7 +40,7 @@ from ..core.ops.builtin import BackwardGraph, OpDef from ..core.ops.special import Const from ..core.tensor import megbrain_graph as G from ..core.tensor.utils import setscalar -from ..utils.naming import auto_naming +from ..utils.naming import AutoNaming from .sublinear_memory_config import SublinearMemoryConfig @@ -297,9 +297,7 @@ class trace: 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() - name = ( - auto_naming.get_scope() + "." + (x.c_name if x.c_name else x._name) - ) + name = AutoNaming.gen_name(x) info.name = name info.external = True info.device = x.device @@ -845,17 +843,17 @@ class trace: ivars.append(h2v[h]) ovars = G.apply_normal_varnode(op, *ivars) - auto_naming.record_opnode(ovars[0].op) + AutoNaming.record_opnode(ovars[0].op) assert len(ovars) == len(ohandles) h2v.update(zip(ohandles, ovars)) for i in ohandles: - name = auto_naming.get_var_name(i) + name = AutoNaming.get_var_name(i) if name is not None: h2v[i].name = name - auto_naming.remove_duplicate_names() + AutoNaming.remove_duplicate_names() dest_vars = [] for i, h in enumerate(self._output_bindings): @@ -1173,7 +1171,7 @@ def apply_const_compiled_mode(value, dtype, device, is_const, no_cache, name): def apply_with_tracing(op: OpDef, *args: RawTensor): if hasattr(op, "scope"): - op.scope = auto_naming.get_scope() + op.scope = AutoNaming.get_scope() if active_trace._symbolic: outputs = apply_symbolic_mode(op, *args) else: diff --git a/imperative/python/megengine/module/module.py b/imperative/python/megengine/module/module.py index 51e4b503b..cd0af84b8 100644 --- a/imperative/python/megengine/module/module.py +++ b/imperative/python/megengine/module/module.py @@ -16,7 +16,7 @@ from ..logger import get_logger from ..tensor import Parameter, Tensor from ..utils.deprecation import deprecated from ..utils.hook import HookHandler -from ..utils.naming import auto_naming +from ..utils.naming import AutoNaming logger = get_logger(__name__) @@ -111,7 +111,7 @@ class Module(metaclass=ABCMeta): self._forward_hooks = OrderedDict() # used for profiler and automatic naming - self._name = "{anonymous}" + self._name = None @abstractmethod def forward(self, inputs): @@ -137,7 +137,7 @@ class Module(metaclass=ABCMeta): return HookHandler(self._forward_hooks, hook) def __call__(self, *inputs, **kwargs): - auto_naming.push_scope(self.name if self.name is not None else self._name) + AutoNaming.push_scope(self.name if self.name is not None else self._name) for hook in self._forward_pre_hooks.values(): modified_inputs = hook(self, inputs) if modified_inputs is not None: @@ -151,7 +151,7 @@ class Module(metaclass=ABCMeta): modified_outputs = hook(self, inputs, outputs) if modified_outputs is not None: outputs = modified_outputs - auto_naming.pop_scope() + AutoNaming.pop_scope() return outputs def _flatten( diff --git a/imperative/python/megengine/tensor.py b/imperative/python/megengine/tensor.py index de4afd33d..c343e2f09 100644 --- a/imperative/python/megengine/tensor.py +++ b/imperative/python/megengine/tensor.py @@ -20,7 +20,7 @@ from .core.tensor.array_method import ArrayMethodMixin from .device import _valid_device, get_default_device from .logger import get_logger from .utils.deprecation import deprecated -from .utils.naming import auto_naming +from .utils.naming import AutoNaming logger = get_logger(__name__) @@ -168,7 +168,7 @@ class Tensor(_Tensor, ArrayMethodMixin): @name.setter def name(self, name): self.c_name = name - auto_naming.record_var_name(self._mixin_handle, name) + AutoNaming.record_var_name(self._mixin_handle, name) @deprecated(version="1.0", reason="no need to reuse an existing tensor since 1.0") def set_value(self, value): diff --git a/imperative/python/megengine/utils/naming.py b/imperative/python/megengine/utils/naming.py index f5eb71cf7..4a7eed683 100644 --- a/imperative/python/megengine/utils/naming.py +++ b/imperative/python/megengine/utils/naming.py @@ -15,40 +15,57 @@ class AutoNaming: renamed by the user. """ - def __init__(self): - self.scopes = [] - self.c_ops = [] - self.name2ops = {} - self.handle2names = {} + scopes = [] + c_ops = [] + name2ops = {} + handle2names = {} + __cls_attributes__ = {"scopes", "c_ops", "name2ops", "handle2names"} - def clear(self): - for var in vars(self).values(): - var.clear() + @classmethod + def clear(cls): + for attr in cls.__cls_attributes__: + getattr(cls, attr).clear() - def push_scope(self, scope): - push_scope(scope) - self.scopes.append(scope) + @classmethod + def push_scope(cls, scope): + if scope is not None: + push_scope(scope) + cls.scopes.append(scope) - def pop_scope(self): - scope = self.scopes.pop() - pop_scope(scope) + @classmethod + def pop_scope(cls): + scope = cls.scopes.pop() + if scope is not None: + pop_scope(scope) - def get_scope(self): - return ".".join(self.scopes) + @classmethod + def get_scope(cls): + return ".".join(s for s in cls.scopes if s is not None) - def record_var_name(self, handle, name): - self.handle2names[handle] = name + @classmethod + def gen_name(cls, x) -> str: + scope = cls.get_scope() + name = x.c_name if x.c_name else x._name + return scope + "." + name if len(scope) else name - def get_var_name(self, handle): - return self.handle2names.pop(handle, None) + @classmethod + def record_var_name(cls, handle, name): + cls.handle2names[handle] = name - def record_opnode(self, op): - ops = self.name2ops.get(op.name, []) - ops.append(op) - self.name2ops[op.name] = ops + @classmethod + def get_var_name(cls, handle): + return cls.handle2names.pop(handle, None) - def remove_duplicate_names(self): - for key, ops in self.name2ops.items(): + @classmethod + def record_opnode(cls, op): + ops = cls.name2ops.get(op.name, []) + if op not in ops: + ops.append(op) + cls.name2ops[op.name] = ops + + @classmethod + def remove_duplicate_names(cls): + for key, ops in cls.name2ops.items(): if len(ops) == 1: continue for i, op in enumerate(ops): @@ -57,7 +74,4 @@ class AutoNaming: continue for var in op.outputs: var.name = var.name.replace(key, op.name) - self.name2ops.clear() - - -auto_naming = AutoNaming() + cls.name2ops.clear() diff --git a/imperative/python/test/unit/jit/test_tracing.py b/imperative/python/test/unit/jit/test_tracing.py index 3a699441d..03bd1dca4 100644 --- a/imperative/python/test/unit/jit/test_tracing.py +++ b/imperative/python/test/unit/jit/test_tracing.py @@ -28,7 +28,7 @@ from megengine.functional import exp, log from megengine.jit import exclude_from_trace, trace from megengine.module import Module from megengine.random import normal, uniform -from megengine.utils.naming import auto_naming +from megengine.utils.naming import AutoNaming @pytest.mark.parametrize("trace_mode", [False, True]) @@ -141,7 +141,7 @@ def test_dump(): return a + b # prevent from remaining scope from exception test - auto_naming.clear() + AutoNaming.clear() a = tensor([2]) b = tensor([4]) y = f(a, b).numpy() diff --git a/imperative/python/test/unit/utils/test_dump_naming.py b/imperative/python/test/unit/utils/test_dump_naming.py index d6219c38e..44845a0b2 100644 --- a/imperative/python/test/unit/utils/test_dump_naming.py +++ b/imperative/python/test/unit/utils/test_dump_naming.py @@ -18,11 +18,11 @@ from megengine import Parameter, Tensor from megengine.core.tensor import megbrain_graph as G from megengine.jit.tracing import trace from megengine.quantization.quantize import quantize, quantize_qat -from megengine.utils.naming import auto_naming +from megengine.utils.naming import AutoNaming def _dump_and_load(func, symbolic, keep_opr_name=True): - auto_naming.clear() + AutoNaming.clear() func = trace(func, symbolic=symbolic, capture_as_const=True) x = Tensor(np.ones(shape=(2, 3))) func(x).numpy() @@ -103,6 +103,18 @@ def test_without_module(symbolic): assert op.name == "MUL" +@pytest.mark.parametrize("symbolic", [False, True]) +def test_ignore_top_module(symbolic): + class Simple(M.Module): + def forward(self, x): + return x + x + + m = Simple() + op = _dump_and_load(m, symbolic)[-1] + assert op.name == "ADD" + assert op.outputs[0].name == "ADD" + + @pytest.mark.parametrize("symbolic", [False, True]) def test_with_submodule(symbolic): class Simple(M.Module): @@ -196,7 +208,7 @@ def test_not_keep_opr_name(): return 2 * x op = _dump_and_load(f, True, False)[-1] - assert op.name == "MUL(x,2[2])[4]" + assert op.name == "MUL(x,const<2>[2])[4]" @pytest.mark.parametrize("symbolic", [False, True]) diff --git a/src/opr/impl/io.cpp b/src/opr/impl/io.cpp index e8753be63..0d9fa34c5 100644 --- a/src/opr/impl/io.cpp +++ b/src/opr/impl/io.cpp @@ -419,7 +419,7 @@ void ImmutableTensor::Value::setup(CompNode cn, const HostTensorND &val) { if (one_elem(val.shape())) { float v; static_cast_dtype(&v, val.dtype(), val.raw_ptr()); - m_summary = ssprintf("%.3g", v); + m_summary = ssprintf("const<%.3g>", v); if (val.shape().ndim != 1) { m_summary += val.shape().to_string(); } -- GitLab