Function quantize_qat

imperative/python/megengine/quantization/quantize.py:84–135

r"""Recursively convert float :class:`~.Module` to :class:`~.QATModule` through :meth:`~.Module.apply` and set qconfig relatively. Args: module: root module to do convert recursively. inplace: whether to convert submodules in-place. qconfig: an instance of :class:`~.

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

def quantize_qat(
    module: Module,
    inplace: bool = True,
    qconfig: QConfig = ema_fakequant_qconfig,
    mapping: dict = None,
):
    r"""Recursively convert float :class:`~.Module` to :class:`~.QATModule`
    through :meth:`~.Module.apply` and set qconfig relatively.

    Args:
        module: root module to do convert recursively.
        inplace: whether to convert submodules in-place.
        qconfig: an instance of :class:`~.QConfig` to be set as submodules' quant config. Default: ``ema_fakequant_qconfig``.
        mapping: a dict indicating how to convert custom modules from Module to QATModule.
            Will be combined with internal default convert mapping dict.

    Returns:
        Return type: Module. Quantized module.

    Examples:
        >>> import megengine.quantization as Q
        >>> model = Net()                                           # doctest: +SKIP
        >>> Q.quantize_qat(model, qconfig=Q.ema_fakequant_qconfig)  # doctest: +SKIP
    """
    assert isinstance(inplace, bool) and isinstance(
        qconfig, QConfig
    ), "the type of input args(inplace and qconfig) is wrong, please check them."

    if not inplace:
        module = deepcopy(module)

    convert_dict = copy(_float2qat_dict)
    if mapping is not None:
        convert_dict.update(mapping)
    quantable_modules = tuple(convert_dict.keys())

    def is_quantable(mod: Module):
        return isinstance(mod, quantable_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_quantable)
    ):
        # only convert top quantable module.
        if is_quantable(parent) or submodule.quantize_disabled:
            continue

        new_mod = convert_dict[type(submodule)].from_float_module(submodule)
        set_expand_structure(module, key, new_mod)

    propagate_qconfig(module, qconfig)
    return module


def reset_qconfig(module: Module, qconfig: QConfig, inplace: bool = True):