提交 888c7f1f 编写于 作者: M Megvii Engine Team

fix(mge/module): fix auto-naming error when there are containers in the module

GitOrigin-RevId: 6b8f0433216ad0a9454b644d2e9b0989c87836bb
上级 d13b6a4a
......@@ -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):
......
......@@ -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
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册