MegEngine/python_module/megengine/quantization/qconfig.py

# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
from functools import partial

from ..module import Module
from .fake_quant import FakeQuantize
from .observer import ExponentialMovingAverageObserver, MinMaxObserver


class QConfig:
    """
    A config class indicating how to do quantize toward :class:`~.QATModule`'s
    ``activation`` and ``weight``.

    And ``fake_quant`` parameter to indicate

    See :meth:`~.QATModule.set_qconfig` for detail usage.

    :param weight_observer: interface to instantiate an :class:`~.Observer` indicating
-        how to collect scales and zero_point of wegiht.
    :param act_observer: similar to ``weight_observer`` but toward activation.
    :param fake_quant: interface to instantiate a :class:`~.FakeQuantize` indicating
        how to do fake_quant calculation. can be invoked multi times to get different
        instance for each target tensor, for better control on enable and disable.

    Examples:

    .. code-block::

        # Default EMA QConfig for QAT.
        ema_fakequant_qconfig = QConfig(
            weight_observer=MinMaxObserver,
            act_observer=ExponentialMovingAverageObserver,
            fake_quant=FakeQuantize,
        )
    """

    def __init__(
        self, act_observer, weight_observer, fake_quant,
    ):
        if isinstance(act_observer, Module) or isinstance(weight_observer, Module):
            raise ValueError(
                "QConfig must not receive observer instance, please pass observer"
                " class generator using `partial(Observer, ...)` instead. Use"
                " partial(MyObserver, x=1) to override arguments to constructor if needed"
            )
        self.act_observer = act_observer
        self.weight_observer = weight_observer
        self.fake_quant = fake_quant


# Default QAT QConfigs
min_max_fakequant_qconfig = QConfig(
    weight_observer=MinMaxObserver,
    act_observer=MinMaxObserver,
    fake_quant=FakeQuantize,
)

ema_fakequant_qconfig = QConfig(
    weight_observer=MinMaxObserver,
    act_observer=ExponentialMovingAverageObserver,
    fake_quant=FakeQuantize,
)