Function quantize

imperative/python/megengine/quantization/quantize.py:52–81

r"""Recursively convert :class:`~.QATModule` to :class:`~.QuantizedModule` through :meth:`~.Module.apply`. Args: module: root module to do convert recursively. inplace: whether to convert submodules in-place. mapping: a dict indicating how to convert custom modules f

(module: Module, inplace: bool = True, mapping: dict = None)

def quantize(module: Module, inplace: bool = True, mapping: dict = None):
    r"""Recursively convert :class:`~.QATModule` to :class:`~.QuantizedModule`
    through :meth:`~.Module.apply`.

    Args:
        module: root module to do convert recursively.
        inplace: whether to convert submodules in-place.
        mapping: a dict indicating how to convert custom modules from QATModule to
            QuantizedModule. Will be combined with internal default convert mapping dict.
    """

    if not inplace:
        module = deepcopy(module)

    convert_dict = copy(_qat2quantized_dict)
    if mapping is not None:
        convert_dict.update(mapping)
    qat_modules = tuple(convert_dict.keys())

    def is_qat(mod: Module):
        return isinstance(mod, qat_modules)

    # must use list to avoid replacement influencing successor modules
    for key, submodule, parent in list(
        module._flatten(with_key=True, with_parent=True, predicate=is_qat)
    ):
        new_mod = convert_dict[type(submodule)].from_qat_module(submodule)
        set_expand_structure(module, key, new_mod)

    return module


def quantize_qat(