From 48a7d333c3dfcf0ae1569417639e363aadb28c2a Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Wed, 24 Mar 2021 11:34:21 +0800 Subject: [PATCH] fix(mge/module): fix auto-naming error when there are containers in the module GitOrigin-RevId: 6b8f0433216ad0a9454b644d2e9b0989c87836bb --- imperative/python/megengine/module/module.py | 6 ++-- .../test/unit/utils/test_dump_naming.py | 31 ++++++++++++++++++- 2 files changed, 33 insertions(+), 4 deletions(-) diff --git a/imperative/python/megengine/module/module.py b/imperative/python/megengine/module/module.py index aab4e87dd..51e4b503b 100644 --- a/imperative/python/megengine/module/module.py +++ b/imperative/python/megengine/module/module.py @@ -607,10 +607,10 @@ class Module(metaclass=ABCMeta): def __getattribute__(self, name: str): value = super().__getattribute__(name) - if name == "_name": + if name == "__dict__": return value - if isinstance(value, (Tensor, Module)): - value._name = name + for prefix, variable in _expand_structure(name, value): + variable._name = prefix return value def __setattr__(self, name: str, value): diff --git a/imperative/python/test/unit/utils/test_dump_naming.py b/imperative/python/test/unit/utils/test_dump_naming.py index eabab1b44..d6219c38e 100644 --- a/imperative/python/test/unit/utils/test_dump_naming.py +++ b/imperative/python/test/unit/utils/test_dump_naming.py @@ -123,6 +123,35 @@ def test_with_submodule(symbolic): assert ops[-1].outputs[0].name == "simple.linear.ADD" +@pytest.mark.parametrize("symbolic", [False, True]) +def test_with_submodule_in_container(symbolic): + class Simple(M.Module): + def __init__(self, name): + super().__init__() + self.name = name + self.l0 = [M.Linear(3, 3) for _ in range(2)] + self.l1 = tuple(self.l0) + self.l2 = dict(zip(["l2-0", "l2-1"], self.l0)) + + def forward(self, x): + for i in range(2): + x = self.l0[i](x) + x = self.l1[i](x) + x = self.l2["l2-%d" % i](x) + return x + + m = Simple("simple") + + ops = _dump_and_load(m, symbolic) + assert ops[-1].outputs[0].name == "simple.l2.l2-1.ADD" + assert ops[-1].name == "simple.l2.l2-1.ADD" + assert ops[-2].name == "simple.l2.l2-1.MatrixMul" + assert ops[-3].name == "simple.l1.1.ADD" + assert ops[-4].name == "simple.l1.1.MatrixMul" + assert ops[-5].name == "simple.l0.1.ADD" + assert ops[-6].name == "simple.l0.1.MatrixMul" + + @pytest.mark.parametrize("symbolic", [False, True]) def test_named_submodule(symbolic): class Simple(M.Module): @@ -264,4 +293,4 @@ def test_quantized_module_user_naming_param(symbolic): (matrix_mul_op,) = [op for op in ops if op.name == "simple.linear.MatrixMul"] for var in matrix_mul_op.inputs: assert var.name in ("simple.quant.TypeCvt", "simple.linear.user-weight") - # BUG bias' name does not meet expectations because of astype operator after quantization + # WONTFIX: bias' name does not meet expectations because of astype operator after quantization -- GitLab