提交 15d1bd4d 编写于 作者: M Megvii Engine Team 提交者: Xu Xinran

fix(mge/quantization): use list rather than generator to modify submodules

GitOrigin-RevId: b5a3851330e7ffeb2aa9ad1a42e8ade8a0817506
上级 7b2c5a73
...@@ -65,9 +65,9 @@ def quantize(module: Module, inplace=True): ...@@ -65,9 +65,9 @@ def quantize(module: Module, inplace=True):
def is_qat(mod: Module): def is_qat(mod: Module):
return isinstance(mod, qat_modules) return isinstance(mod, qat_modules)
# no need to pass prefix and get pure key of parent Module. # must use list to avoid replacement influencing successor modules
for key, submodule, parent in module._flatten( for key, submodule, parent in list(
with_key=True, with_parent=True, predicate=is_qat module._flatten(with_key=True, with_parent=True, predicate=is_qat)
): ):
new_mod = _qat2quantized_dict[type(submodule)].from_qat_module(submodule) new_mod = _qat2quantized_dict[type(submodule)].from_qat_module(submodule)
if isinstance(parent, Float.Sequential): if isinstance(parent, Float.Sequential):
...@@ -100,9 +100,9 @@ def quantize_qat( ...@@ -100,9 +100,9 @@ def quantize_qat(
def is_quantable(mod: Module): def is_quantable(mod: Module):
return isinstance(mod, quantable_modules) return isinstance(mod, quantable_modules)
# no need to pass prefix and get pure key of parent Module. # must use list to avoid replacement influencing successor modules
for key, submodule, parent in module._flatten( for key, submodule, parent in list(
with_key=True, with_parent=True, predicate=is_quantable module._flatten(with_key=True, with_parent=True, predicate=is_quantable)
): ):
# only convert top quantable module. # only convert top quantable module.
if is_quantable(parent): if is_quantable(parent):
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册