chore: import upstream snapshot with attribution

python/paddle/quantization/__init__.py

@@ -0,0 +1,53 @@
+"""Quantization Module"""
+
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .base_observer import BaseObserver
+from .base_quanter import BaseQuanter
+from .config import QuantConfig
+from .factory import quanter
+from .imperative.ptq import (  # noqa: F401
+    ImperativePTQ,
+)
+from .imperative.ptq_config import (  # noqa: F401
+    PTQConfig,
+    default_ptq_config,
+)
+from .imperative.ptq_quantizer import (  # noqa: F401
+    SUPPORT_ACT_QUANTIZERS,
+    SUPPORT_WT_QUANTIZERS,
+    AbsmaxQuantizer,
+    BaseQuantizer,
+    HistQuantizer,
+    KLQuantizer,
+    PerChannelAbsmaxQuantizer,
+)
+from .imperative.ptq_registry import (  # noqa: F401
+    PTQRegistry,
+)
+from .imperative.qat import (  # noqa: F401
+    ImperativeQuantAware,
+)
+from .ptq import PTQ
+from .qat import QAT
+
+__all__ = [
+    "QuantConfig",
+    "BaseQuanter",
+    "BaseObserver",
+    "quanter",
+    "QAT",
+    "PTQ",
+]

python/paddle/quantization/base_observer.py

@@ -0,0 +1,34 @@
+"""Abstract observer class."""
+
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+import abc
+
+from .base_quanter import BaseQuanter
+
+
+class BaseObserver(BaseQuanter, metaclass=abc.ABCMeta):
+    r"""
+    Built-in observers and customized observers should extend this base observer
+    and implement abstract methods.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    @abc.abstractmethod
+    def cal_thresholds(self) -> None:
+        pass

python/paddle/quantization/base_quanter.py

@@ -0,0 +1,70 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+import abc
+from typing import TYPE_CHECKING, Any
+
+from paddle.nn import Layer
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
+    import numpy.typing as npt
+
+    from paddle import Tensor
+
+
+class BaseQuanter(Layer, metaclass=abc.ABCMeta):
+    r"""
+    Built-in quanters and customized quanters should extend this base quanter
+    and implement abstract methods.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    @abc.abstractmethod
+    def forward(self, input: Tensor) -> Tensor | npt.NDArray[Any]:
+        pass
+
+    @abc.abstractmethod
+    def scales(self) -> Tensor | npt.NDArray[Any]:
+        r"""
+        Get the scales used for quantization.
+        It can be none which means the quanter didn't hold scales for quantization.
+        """
+        pass
+
+    @abc.abstractmethod
+    def zero_points(self) -> Tensor | npt.NDArray[Any]:
+        r"""
+        Get the zero points used for quantization.
+        It can be none which means the quanter didn't hold zero points for quantization.
+        """
+        pass
+
+    @abc.abstractmethod
+    def quant_axis(self) -> int | Iterable[int]:
+        r"""
+        Get the axis of quantization. None means tensor-wise quantization.
+        """
+        pass
+
+    @abc.abstractmethod
+    def bit_length(self) -> int | Iterable[int]:
+        r"""
+        Get the bit length of quantization.
+        """
+        pass

python/paddle/quantization/config.py

@@ -0,0 +1,497 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+import copy
+from typing import TYPE_CHECKING
+
+import paddle
+from paddle import nn
+
+from .wrapper import ObserveWrapper
+
+if TYPE_CHECKING:
+    from paddle.nn import Layer
+
+    from .factory import QuanterFactory
+
+
+# TODO: Implement quanted layer and fill the mapping dict
+DEFAULT_QAT_LAYER_MAPPINGS: dict[Layer, Layer] = {
+    nn.quant.Stub: nn.quant.stub.QuanterStub,
+    nn.Linear: nn.quant.qat.QuantedLinear,
+    nn.Conv2D: nn.quant.qat.QuantedConv2D,
+}
+
+DEFAULT_LEAVES = [nn.ReLU, nn.AvgPool2D]
+
+
+class SingleLayerConfig:
+    r"""
+    Configure how to quantize the activations and weights of a single layer.
+
+    Args:
+        activation(QuanterFactory): The factory to create instance of quanter used to quantize activations.
+        weight(QuanterFactory): The factory to create instance of quanter used to quantize weights.
+    """
+
+    def __init__(
+        self, activation: QuanterFactory, weight: QuanterFactory
+    ) -> None:
+        self._activation = activation
+        self._weight = weight
+
+    @property
+    def activation(self) -> QuanterFactory:
+        return self._activation
+
+    @property
+    def weight(self) -> QuanterFactory:
+        return self._weight
+
+    def __str__(self):
+        return f"activation: {self._activation}\nweight: {self._weight}"
+
+
+class QuantConfig:
+    r"""
+    Configure how to quantize a model or a part of the model. It will map each layer to
+    an instance of SingleLayerConfig by the settings. It provides diverse methods to set
+    the strategies of quantization.
+
+    Args:
+        activation(QuanterFactory | None): The global quantizer used to quantize the activations.
+        weight(QuanterFactory | None): The global quantizer used to quantize the weights.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> from paddle.quantization import QuantConfig
+            >>> from paddle.quantization.quanters import FakeQuanterWithAbsMaxObserver
+
+            >>> quanter = FakeQuanterWithAbsMaxObserver(moving_rate=0.9)
+            >>> q_config = QuantConfig(activation=quanter, weight=quanter)
+            >>> print(q_config)
+            Global config:
+            activation: FakeQuanterWithAbsMaxObserver(name=None,moving_rate=0.9,bit_length=8,dtype=float32)
+            weight: FakeQuanterWithAbsMaxObserver(name=None,moving_rate=0.9,bit_length=8,dtype=float32)
+
+    """
+
+    def __init__(
+        self, activation: QuanterFactory | None, weight: QuanterFactory | None
+    ) -> None:
+        if activation is None and weight is None:
+            self._global_config = None
+        else:
+            self._global_config = SingleLayerConfig(activation, weight)
+        self._layer2config = {}
+        self._prefix2config = {}
+        self._type2config = {}
+        self._model = None
+        self._qat_layer_mapping = copy.deepcopy(DEFAULT_QAT_LAYER_MAPPINGS)
+        self._customized_qat_layer_mapping = {}
+
+        self._customized_leaves = []
+
+    def add_layer_config(
+        self,
+        layer: Layer | list[Layer],
+        activation: QuanterFactory | None = None,
+        weight: QuanterFactory | None = None,
+    ) -> None:
+        r"""
+        Set the quantization config by layer. It has the highest priority among
+        all the setting methods.
+
+        Args:
+            layer(Layer|list[Layer]]): One or a list of layers.
+            activation(QuanterFactory | None): Quanter used for activations. Default is None.
+            weight(QuanterFactory | None): Quanter used for weights. Default is None.
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle
+                >>> from paddle.nn import Linear
+                >>> from paddle.quantization import QuantConfig
+                >>> from paddle.quantization.quanters import FakeQuanterWithAbsMaxObserver
+
+                >>> class Model(paddle.nn.Layer):
+                ...     def __init__(self):
+                ...         super().__init__()
+                ...         self.fc = Linear(576, 120)
+                >>> model = Model()
+                >>> quanter = FakeQuanterWithAbsMaxObserver(moving_rate=0.9)
+                >>> q_config = QuantConfig(activation=None, weight=None)
+                >>> q_config.add_layer_config([model.fc], activation=quanter, weight=quanter)
+                >>> # doctest: +SKIP('random memory address')
+                >>> print(q_config)
+                Global config:
+                None
+                Layer prefix config:
+                {'linear_0': <paddle.quantization.config.SingleLayerConfig object at 0x7fe41a680ee0>}
+
+        """
+        if isinstance(layer, list):
+            for _element in layer:
+                self.add_layer_config(
+                    _element, activation=activation, weight=weight
+                )
+        else:
+            self.add_name_config(
+                layer.full_name(), activation=activation, weight=weight
+            )
+
+    def add_name_config(
+        self,
+        layer_name: str | list[str],
+        activation: QuanterFactory | None = None,
+        weight: QuanterFactory | None = None,
+    ) -> None:
+        r"""
+        Set the quantization config by full name of layer. Its priority is
+        lower than `add_layer_config`.
+
+        Args:
+            layer_name(str|list[str]): One or a list of layers' full name.
+            activation(QuanterFactory | None): Quanter used for activations. Default is None.
+            weight(QuanterFactory | None): Quanter used for weights. Default is None.
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle
+                >>> from paddle.nn import Linear
+                >>> from paddle.quantization import QuantConfig
+                >>> from paddle.quantization.quanters import FakeQuanterWithAbsMaxObserver
+
+                >>> class Model(paddle.nn.Layer):
+                ...     def __init__(self):
+                ...         super().__init__()
+                ...         self.fc = Linear(576, 120)
+                >>> model = Model()
+                >>> quanter = FakeQuanterWithAbsMaxObserver(moving_rate=0.9)
+                >>> q_config = QuantConfig(activation=None, weight=None)
+                >>> q_config.add_name_config([model.fc.full_name()], activation=quanter, weight=quanter)
+                >>> # doctest: +SKIP('random memory address')
+                >>> print(q_config)
+                Global config:
+                None
+                Layer prefix config:
+                {'linear_0': <paddle.quantization.config.SingleLayerConfig object at 0x7fe41a680fd0>}
+
+        """
+        if isinstance(layer_name, str):
+            config = SingleLayerConfig(activation, weight)
+            self._prefix2config[layer_name] = config
+        if isinstance(layer_name, list):
+            for _element in layer_name:
+                self.add_name_config(
+                    _element, activation=activation, weight=weight
+                )
+
+    def add_type_config(
+        self,
+        layer_type: type[Layer] | list[type[Layer]],
+        activation: QuanterFactory | None = None,
+        weight: QuanterFactory | None = None,
+    ) -> None:
+        r"""
+        Set the quantization config by the type of layer. The `layer_type` should be
+        subclass of `paddle.nn.Layer`. Its priority is lower than `add_layer_config`
+        and `add_name_config`.
+
+        Args:
+            layer_type(type[Layer] | list[type[Layer]]): One or a list of layers' type. It should be subclass of
+            `paddle.nn.Layer`. Python built-in function `type()` can be used to get the type of a layer.
+            activation(QuanterFactory | None): Quanter used for activations. Default is None.
+            weight(QuanterFactory | None): Quanter used for weights. Default is None.
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle
+                >>> from paddle.nn import Linear
+                >>> from paddle.quantization import QuantConfig
+                >>> from paddle.quantization.quanters import (
+                ...     FakeQuanterWithAbsMaxObserver,
+                ... )
+
+                >>> class Model(paddle.nn.Layer):
+                ...     def __init__(self):
+                ...         super().__init__()
+                ...         self.fc = Linear(576, 120)
+                >>> model = Model()
+                >>> quanter = FakeQuanterWithAbsMaxObserver(moving_rate=0.9)
+                >>> q_config = QuantConfig(activation=None, weight=None)
+                >>> q_config.add_type_config([Linear], activation=quanter, weight=quanter)
+                >>> # doctest: +SKIP('random memory address')
+                >>> print(q_config)
+                Global config:
+                None
+                Layer type config:
+                {<class 'paddle.nn.layer.common.Linear'>: <paddle.quantization.config.SingleLayerConfig object at 0x7fe41a680a60>}
+
+        """
+        if isinstance(layer_type, type) and issubclass(
+            layer_type, paddle.nn.Layer
+        ):
+            config = SingleLayerConfig(activation, weight)
+            self._type2config[layer_type] = config
+        if isinstance(layer_type, list):
+            for _element in layer_type:
+                self.add_type_config(
+                    _element, activation=activation, weight=weight
+                )
+
+    def add_qat_layer_mapping(
+        self, source: type[Layer], target: type[Layer]
+    ) -> None:
+        r"""
+        Add rules converting layers to simulated quantization layers
+        before quantization-aware training. It will convert layers
+        with type `source` to layers with type `target`. `source` and
+        `target` should be subclass of `paddle.nn.Layer`. And a default
+        mapping is provided by property `default_qat_layer_mapping`.
+
+        Args:
+            source(type[Layer]): The type of layers that will be converted.
+            target(type[Layer]): The type of layers that will be converted to.
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle
+                >>> from paddle.nn import Conv2D
+                >>> from paddle.quantization import QuantConfig
+                >>> from paddle.quantization.quanters import (
+                ...     FakeQuanterWithAbsMaxObserver,
+                ... )
+                >>> quanter = FakeQuanterWithAbsMaxObserver(moving_rate=0.9)
+                >>> q_config = QuantConfig(activation=None, weight=None)
+                >>> class CustomizedQuantedConv2D(paddle.nn.Layer):
+                ...     def forward(self, x):
+                ...         pass
+                ...         # add some code for quantization simulation
+                >>> q_config.add_qat_layer_mapping(Conv2D, CustomizedQuantedConv2D)
+        """
+        assert isinstance(source, type) and issubclass(
+            source, paddle.nn.Layer
+        ), (
+            "The source layer to be placed should be a subclass of paddle.nn.Layer"
+        )
+        assert isinstance(target, type) and issubclass(
+            target, paddle.nn.Layer
+        ), "The target layer should be a subclass of paddle.nn.qat.Layer"
+        self._qat_layer_mapping[source] = target
+        self._customized_qat_layer_mapping[source] = target
+
+    def add_customized_leaf(self, layer_type: type[Layer]) -> None:
+        r"""
+        Declare the customized layer as leaf of model for quantization.
+        The leaf layer is quantized as one layer. The sublayers of
+        leaf layer will not be quantized.
+
+        Args:
+            layer_type(type[Layer]): The type of layer to be declared as leaf.
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> from paddle.nn import Sequential
+                >>> from paddle.quantization import QuantConfig
+                >>> from paddle.quantization.quanters import FakeQuanterWithAbsMaxObserver
+                >>> q_config = QuantConfig(activation=None, weight=None)
+                >>> q_config.add_customized_leaf(Sequential)
+
+        """
+        self._customized_leaves.append(layer_type)
+
+    @property
+    def customized_leaves(self) -> list[type[Layer]]:
+        r"""
+        Get all the customized leaves.
+        """
+        return self._customized_leaves
+
+    def _need_observe(self, layer: Layer):
+        r"""
+        Whether the layer should be observed by observer.
+        """
+        return self._is_leaf(layer) and self._has_observer_config(layer)
+
+    def _get_qat_layer(self, layer: Layer):
+        q_config = self._get_config_by_layer(layer)
+
+        target_type = self._customized_qat_layer_mapping.get(
+            type(layer), self.qat_layer_mappings.get(type(layer))
+        )
+        return target_type(layer, q_config)
+
+    def _has_observer_config(self, layer: Layer):
+        r"""
+        Whether the layer has been configured for activation quantization.
+        """
+        _config = self._get_config_by_layer(layer)
+        return _config is not None and _config.activation is not None
+
+    def _is_leaf(self, layer: Layer):
+        return (
+            self._is_default_leaf(layer)
+            or self._is_real_leaf(layer)
+            or self._is_customized_leaf(layer)
+        )
+
+    def _is_default_leaf(self, layer: Layer):
+        return type(layer) in DEFAULT_LEAVES
+
+    def _is_real_leaf(self, layer: Layer):
+        r"""
+        The leaf is real leaf when it has no sublayers.
+        """
+        return layer._sub_layers is None or len(layer._sub_layers) == 0
+
+    def _is_customized_leaf(self, layer: Layer):
+        return type(layer) in self.customized_leaves
+
+    def _get_observer(self, layer: Layer):
+        r"""
+        Create an instance of observer or quanter according to the
+        given layer's quantization config.
+        """
+        _config = self._get_config_by_layer(layer)
+        _observer = None if _config is None else _config.activation
+        return None if _observer is None else _observer._instance(layer)
+
+    def _get_observe_wrapper(self, layer: Layer):
+        _observer = self._get_observer(layer)
+        return ObserveWrapper(_observer, layer)
+
+    @property
+    def qat_layer_mappings(self) -> dict[Layer, Layer]:
+        return self._qat_layer_mapping
+
+    @property
+    def default_qat_layer_mapping(self) -> dict[Layer, Layer]:
+        return DEFAULT_QAT_LAYER_MAPPINGS
+
+    @property
+    def global_config(self) -> SingleLayerConfig:
+        return self._global_config
+
+    def _get_config_by_layer(self, layer) -> SingleLayerConfig:
+        return self._layer2config.get(layer, None)
+
+    def _is_quantifiable(self, layer: Layer):
+        r"""
+        The layer is quantifiable when it configured by activation quanter/observer
+        or weight quanter/observer.
+        """
+        return layer in self._layer2config
+
+    def _specify(self, model: Layer):
+        r"""
+        Specify the quantization config of each sublayer in model.
+        For each layer in sublayers of mode,
+        1. Set the config by global config
+        2. Overwrite the config with parents' config
+        3. Overwrite the config with config set by layer's type
+        4. Overwrite the config with config set by layer's full name
+        5. Overwrite the config with config set by layer
+
+        Args:
+            model(Layer): The model to be specified by the config.
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle
+                >>> from paddle.nn import Linear, Sequential
+                >>> from paddle.quantization import QuantConfig
+                >>> from paddle.quantization.quanters import FakeQuanterWithAbsMaxObserver
+
+                >>> class Model(paddle.nn.Layer):
+                ...     def __init__(self):
+                ...         super().__init__()
+                ...         self.fc = Sequential(Linear(576, 120), Linear(576, 120))
+                >>> model = Model()
+                >>> quanter = FakeQuanterWithAbsMaxObserver(moving_rate=0.9)
+                >>> q_config = QuantConfig(activation=None, weight=None)
+                >>> q_config.add_layer_config([model.fc], activation=quanter, weight=quanter)
+                >>> q_config._specify(model)
+        """
+        self._model = model
+        self._specify_helper(self._model)
+
+    def _specify_helper(self, model: Layer):
+        for child in model.children():
+            layer_prefix = child.full_name()
+            config = self._layer2config.get(model, self.global_config)
+
+            config = self._type2config.get(type(child), config)
+            config = self._prefix2config.get(layer_prefix, config)
+            if config is not None:
+                self._layer2config[child] = config
+            self._specify_helper(child)
+        return self
+
+    def details(self) -> str:
+        r"""
+        Get the formatted details of current config.
+        """
+        if self._model is None:
+            return self.__str__()
+        return self._details_helper(self._model)
+
+    def _details_helper(self, layer: Layer):
+        sublayer_lines = []
+        for name, sublayer in layer.named_children():
+            sublayer_str = self._details_helper(sublayer)
+            sublayer_str = self._addindent(sublayer_str, 2)
+            if sublayer in self._layer2config:
+                sublayer_lines.append(
+                    '('
+                    + name
+                    + '): '
+                    + sublayer_str
+                    + ', '
+                    + str(self._layer2config[sublayer])
+                )
+
+        final_str = layer.__class__.__name__ + '('
+        if sublayer_lines:
+            final_str += '\n  ' + '\n  '.join(sublayer_lines) + '\n'
+
+        final_str += ')'
+        return final_str
+
+    def _addindent(self, string, indent):
+        s1 = string.split('\n')
+        if len(s1) == 1:
+            return string
+        s2 = []
+        for idx, line in enumerate(s1):
+            if idx > 0:
+                s2.append(str((indent * ' ') + line))
+        return s1[0] + '\n' + '\n'.join(s2)
+
+    def __str__(self):
+        result = ""
+        result += f"Global config:\n{self._global_config}\n"
+        if len(self._type2config) > 0:
+            result += f"Layer type config:\n{self._type2config}\n"
+        if len(self._prefix2config) > 0:
+            result += f"Layer prefix config: \n{self._prefix2config}\n"
+        return result

python/paddle/quantization/factory.py

@@ -0,0 +1,142 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+import abc
+import inspect
+from functools import partial
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from paddle.nn import Layer
+
+    from .base_quanter import BaseQuanter
+
+
+class ClassWithArguments(metaclass=abc.ABCMeta):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        self._args = args
+        self._kwargs = kwargs
+
+    @property
+    def args(self):
+        return self._args
+
+    @property
+    def kwargs(self):
+        return self._kwargs
+
+    @abc.abstractmethod
+    def _get_class(self):
+        pass
+
+    def __str__(self):
+        args_str = ",".join(
+            list(self.args) + [f"{k}={v}" for k, v in self.kwargs.items()]
+        )
+        return f"{self.__class__.__name__}({args_str})"
+
+    def __repr__(self):
+        return self.__str__()
+
+
+class QuanterFactory(ClassWithArguments):
+    r"""
+    The factory holds the quanter's class information and
+    the arguments used to create quanter instance.
+    """
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self.partial_class = None
+
+    def _instance(self, layer: Layer) -> BaseQuanter:
+        r"""
+        Create an instance of quanter for target layer.
+        """
+        if self.partial_class is None:
+            self.partial_class = partial(
+                self._get_class(), *self.args, **self.kwargs
+            )
+        return self.partial_class(layer)
+
+
+ObserverFactory = QuanterFactory
+
+
+def quanter(
+    class_name: str,
+) -> Callable[[type[BaseQuanter]], type[BaseQuanter]]:
+    r"""
+    Annotation to declare a factory class for quanter.
+
+    Args:
+        class_name (str): The name of factory class to be declared.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # type: ignore
+            >>> # doctest: +SKIP('need 2 file to run example')
+            >>> # Given codes in ./customized_quanter.py
+            >>> from paddle.quantization import quanter
+            >>> from paddle.quantization import BaseQuanter
+            >>> @quanter("CustomizedQuanter")
+            >>> class CustomizedQuanterLayer(BaseQuanter):
+            ...     def __init__(self, arg1, kwarg1=None):
+            ...         pass
+
+            >>> # Used in ./test.py
+            >>> # from .customized_quanter import CustomizedQuanter
+            >>> from paddle.quantization import QuantConfig
+            >>> arg1_value = "test"
+            >>> kwarg1_value = 20
+            >>> quanter = CustomizedQuanter(arg1_value, kwarg1=kwarg1_value)
+            >>> q_config = QuantConfig(activation=quanter, weight=quanter)
+
+    """
+
+    def wrapper(target_class: type[BaseQuanter]) -> type[BaseQuanter]:
+        init_function_str = f"""
+def init_function(self, *args, **kwargs):
+    super(type(self), self).__init__(*args, **kwargs)
+    import importlib
+    module = importlib.import_module("{target_class.__module__}")
+    my_class = getattr(module, "{target_class.__name__}")
+    globals()["{target_class.__name__}"] = my_class
+def get_class_function(self):
+    return {target_class.__name__}
+locals()["init_function"]=init_function
+locals()["get_class_function"]=get_class_function
+"""
+        exec(init_function_str)
+        frm = inspect.stack()[1]
+        mod = inspect.getmodule(frm[0])
+        new_class = type(
+            class_name,
+            (QuanterFactory,),
+            {
+                "__init__": locals()["init_function"],
+                "_get_class": locals()["get_class_function"],
+            },
+        )
+        setattr(mod, class_name, new_class)
+        if "__all__" in mod.__dict__:
+            mod.__all__.append(class_name)
+
+        return target_class
+
+    return wrapper

python/paddle/quantization/imperative/__init__.py

@@ -0,0 +1,34 @@
+#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from . import (
+    ptq,  # noqa: F401
+    ptq_config,  # noqa: F401
+    ptq_quantizer,  # noqa: F401
+    ptq_registry,  # noqa: F401
+    qat,  # noqa: F401
+)
+from .ptq import ImperativePTQ  # noqa: F401
+from .ptq_config import PTQConfig, default_ptq_config  # noqa: F401
+from .ptq_quantizer import (  # noqa: F401
+    SUPPORT_ACT_QUANTIZERS,
+    SUPPORT_WT_QUANTIZERS,
+    AbsmaxQuantizer,
+    BaseQuantizer,
+    HistQuantizer,
+    KLQuantizer,
+    PerChannelAbsmaxQuantizer,
+)
+from .ptq_registry import PTQRegistry  # noqa: F401
+from .qat import ImperativeQuantAware  # noqa: F401

python/paddle/quantization/imperative/fuse_utils.py

@@ -0,0 +1,221 @@
+#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+
+import paddle
+from paddle import nn
+
+from . import utils
+
+
+class Identity(nn.Layer):
+    '''a layer to replace bn or relu layers'''
+
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+
+    def forward(self, input):
+        return input
+
+
+def fuse_conv_bn(model):
+    is_train = False
+    if model.training:
+        model.eval()
+        is_train = True
+    fuse_list = []
+    tmp_pair = [None, None]
+    for name, layer in model.named_sublayers():
+        if isinstance(layer, nn.Conv2D):
+            tmp_pair[0] = name
+        if isinstance(layer, nn.BatchNorm2D):
+            tmp_pair[1] = name
+
+        if tmp_pair[0] and tmp_pair[1] and len(tmp_pair) == 2:
+            fuse_list.append(tmp_pair)
+            tmp_pair = [None, None]
+    model = fuse_layers(model, fuse_list)
+    if is_train:
+        model.train()
+
+
+def fuse_layers(model, layers_to_fuse, inplace=False):
+    '''
+    fuse layers in layers_to_fuse
+
+    Args:
+        model(paddle.nn.Layer): The model to be fused.
+        layers_to_fuse(list): The layers' names to be fused. For
+            example,"fuse_list = [["conv1", "bn1"], ["conv2", "bn2"]]".
+            A TypeError would be raised if "fuse" was set as
+            True but "fuse_list" was None.
+                              Default: None.
+        inplace(bool): Whether apply fusing to the input model.
+                       Default: False.
+
+    Return
+        fused_model(paddle.nn.Layer): The fused model.
+    '''
+    if inplace is False:
+        model = copy.deepcopy(model)
+    for layers in layers_to_fuse:
+        _fuse_layers(model, layers)
+    return model
+
+
+def _fuse_layers(model, layers_list):
+    '''fuse all the layers in layers_list'''
+    layer_list = []
+    for layer_name in layers_list:
+        parent_layer, sub_name = utils.find_parent_layer_and_sub_name(
+            model, layer_name
+        )
+        layer_list.append(getattr(parent_layer, sub_name))
+    new_layers = _fuse_func(layer_list)
+    for i, item in enumerate(layers_list):
+        parent_layer, sub_name = utils.find_parent_layer_and_sub_name(
+            model, item
+        )
+        setattr(parent_layer, sub_name, new_layers[i])
+
+
+def _fuse_func(layer_list):
+    '''choose the fuse method and fuse layers'''
+    types = tuple(type(m) for m in layer_list)
+    fusion_method = types_to_fusion_method.get(types, None)
+    new_layers = [None] * len(layer_list)
+    fused_layer = fusion_method(*layer_list)
+    for handle_id, pre_hook_fn in layer_list[0]._forward_pre_hooks.items():
+        fused_layer.register_forward_pre_hook(pre_hook_fn)
+        del layer_list[0]._forward_pre_hooks[handle_id]
+    for handle_id, hook_fn in layer_list[-1]._forward_post_hooks.items():
+        fused_layer.register_forward_post_hook(hook_fn)
+        del layer_list[-1]._forward_post_hooks[handle_id]
+    new_layers[0] = fused_layer
+    for i in range(1, len(layer_list)):
+        identity = Identity()
+        identity.training = layer_list[0].training
+        new_layers[i] = identity
+    return new_layers
+
+
+def _fuse_conv_bn(conv, bn):
+    '''fuse conv and bn for train or eval'''
+    assert conv.training == bn.training, (
+        "Conv and BN both must be in the same mode (train or eval)."
+    )
+    if conv.training:
+        assert bn._num_features == conv._out_channels, (
+            'Output channel of Conv2d must match num_features of BatchNorm2d'
+        )
+        raise NotImplementedError
+    else:
+        return _fuse_conv_bn_eval(conv, bn)
+
+
+def _fuse_conv_bn_eval(conv, bn):
+    '''fuse conv and bn for eval'''
+    assert not (conv.training or bn.training), "Fusion only for eval!"
+    fused_conv = copy.deepcopy(conv)
+
+    fused_weight, fused_bias = _fuse_conv_bn_weights(
+        fused_conv.weight,
+        fused_conv.bias,
+        bn._mean,
+        bn._variance,
+        bn._epsilon,
+        bn.weight,
+        bn.bias,
+    )
+    fused_conv.weight.set_value(fused_weight)
+    if fused_conv.bias is None:
+        fused_conv.bias = paddle.create_parameter(
+            shape=[fused_conv._out_channels], is_bias=True, dtype=bn.bias.dtype
+        )
+    fused_conv.bias.set_value(fused_bias)
+    return fused_conv
+
+
+def _fuse_conv_bn_weights(conv_w, conv_b, bn_rm, bn_rv, bn_eps, bn_w, bn_b):
+    '''fuse weights and bias of conv and bn'''
+    if conv_b is None:
+        conv_b = paddle.zeros_like(bn_rm)
+    if bn_w is None:
+        bn_w = paddle.ones_like(bn_rm)
+    if bn_b is None:
+        bn_b = paddle.zeros_like(bn_rm)
+    bn_var_rsqrt = paddle.rsqrt(bn_rv + bn_eps)
+    conv_w = conv_w * (bn_w * bn_var_rsqrt).reshape(
+        [-1] + [1] * (len(conv_w.shape) - 1)
+    )
+    conv_b = (conv_b - bn_rm) * bn_var_rsqrt * bn_w + bn_b
+    return conv_w, conv_b
+
+
+def _fuse_linear_bn(linear, bn):
+    '''fuse linear and bn'''
+    assert linear.training == bn.training, (
+        "Linear and BN both must be in the same mode (train or eval)."
+    )
+    if linear.training:
+        assert bn._num_features == linear.weight.shape[1], (
+            'Output channel of Linear must match num_features of BatchNorm'
+        )
+        raise NotImplementedError
+    else:
+        return _fuse_linear_bn_eval(linear, bn)
+
+
+def _fuse_linear_bn_eval(linear, bn):
+    '''fuse linear and bn for eval'''
+    assert not (linear.training or bn.training), "Fusion only for eval!"
+    fused_linear = copy.deepcopy(linear)
+
+    fused_weight, fused_bias = _fuse_linear_bn_weights(
+        fused_linear.weight,
+        fused_linear.bias,
+        bn._mean,
+        bn._variance,
+        bn._epsilon,
+        bn.weight,
+        bn.bias,
+    )
+    fused_linear.weight.set_value(fused_weight)
+    if fused_linear.bias is None:
+        fused_linear.bias = paddle.create_parameter(
+            shape=[fused_linear.weight.shape[1]],
+            is_bias=True,
+            dtype=bn.bias.dtype,
+        )
+    fused_linear.bias.set_value(fused_bias)
+    return fused_linear
+
+
+def _fuse_linear_bn_weights(
+    linear_w, linear_b, bn_rm, bn_rv, bn_eps, bn_w, bn_b
+):
+    '''fuse weights and bias of linear and bn'''
+    if linear_b is None:
+        linear_b = paddle.zeros_like(bn_rm)
+    bn_scale = bn_w * paddle.rsqrt(bn_rv + bn_eps)
+    fused_w = linear_w * bn_scale.unsqueeze(-1)
+    fused_b = (linear_b - bn_rm) * bn_scale + bn_b
+    return fused_w, fused_b
+
+
+types_to_fusion_method = {
+    (nn.Conv2D, nn.BatchNorm2D): _fuse_conv_bn,
+    (nn.Linear, nn.BatchNorm1D): _fuse_linear_bn,
+}

python/paddle/quantization/imperative/ptq.py

@@ -0,0 +1,485 @@
+#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+import logging
+import os
+
+import numpy as np
+
+import paddle
+from paddle.nn.quant import quant_layers
+
+from ...static.log_helper import get_logger
+from ...static.quantization.utils import (
+    _get_input_name_index,
+    _get_op_input_var_names,
+    _get_op_output_var_names,
+    _get_output_name_index,
+)
+from . import fuse_utils, ptq_config, ptq_hooks, ptq_quantizer, utils
+from .ptq_registry import PTQRegistry
+
+INFER_MODEL_SUFFIX = ".pdmodel"
+INFER_PARAMS_SUFFIX = ".pdiparams"
+
+_logger = get_logger(
+    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
+)
+
+
+class ImperativePTQ:
+    """
+    Static post training quantization.
+    """
+
+    def __init__(self, quant_config=ptq_config.default_ptq_config):
+        """
+        Constructor.
+
+        Args:
+            quant_config(PTQConfig): the config of post training quantization.
+                The config has weight_quantizer and activation_quantizer.
+                In default, the weight_quantizer is PerChannelAbsmaxQuantizer
+                and the activation_quantizer is KLQuantizer.
+        """
+        super().__init__()
+
+        assert isinstance(quant_config, ptq_config.PTQConfig)
+
+        self._quant_config = quant_config
+
+    def quantize(self, model, inplace=False, fuse=False, fuse_list=None):
+        """
+        Add quant config and hook to the target layer.
+
+        Args:
+            model(paddle.nn.Layer): The model to be quantized.
+            inplace(bool): Whether apply quantization to the input model.
+                           Default: False.
+            fuse(bool): Whether to fuse layers.
+                        Default: False.
+            fuse_list(list): The layers' names to be fused. For example,
+                "fuse_list = [["conv1", "bn1"], ["conv2", "bn2"]]".
+                A TypeError would be raised if "fuse" was set as
+                True but "fuse_list" was None.
+                Default: None.
+        Return
+            quantized_model(paddle.nn.Layer): The quantized model.
+        """
+        assert isinstance(model, paddle.nn.Layer), (
+            "The model must be the instance of paddle.nn.Layer."
+        )
+        if not inplace:
+            model = copy.deepcopy(model)
+        if fuse:
+            model.eval()
+            model = fuse_utils.fuse_layers(model, fuse_list)
+        for name, layer in model.named_sublayers():
+            if (
+                PTQRegistry.is_supported_layer(layer)
+                and utils.is_leaf_layer(layer)
+                and not self._is_skip_layer(layer)
+            ):
+                # Add quant config
+                quant_config = copy.deepcopy(self._quant_config)
+                if PTQRegistry.is_simulated_quant_layer(layer):
+                    quant_config.enable_in_act_quantizer = True
+                layer._quant_config = quant_config
+
+                # register hook
+                hook = ptq_hooks.quant_forward_post_hook
+                quant_hook_handle = layer.register_forward_post_hook(hook)
+                quant_config.quant_hook_handle = quant_hook_handle
+                layer._forward_post_hooks.move_to_end(
+                    quant_hook_handle._hook_id, last=False
+                )
+
+        return model
+
+    def save_quantized_model(self, model, path, input_spec=None, **config):
+        """
+        1. Convert the quantized model
+        2. Call jit.save to save the inference model
+        3. Post process the inference model.
+
+        Args:
+            model (Layer): The model to be saved.
+            path (str): The path prefix to save model. The format is
+                ``dirname/file_prefix`` or ``file_prefix``.
+            input_spec (list[InputSpec|Tensor], optional): Describes the input
+                of the saved model's forward method, which can be described by
+                InputSpec or example Tensor. If None, all input variables of
+                the original Layer's forward method would be the inputs of
+                the saved model. Default None.
+            **config (dict, optional): Other save configuration options for
+                compatibility. We do not recommend using these configurations,
+                they may be removed in the future. If not necessary, DO NOT use
+                them. Default None.
+                The following options are currently supported:
+                (1) output_spec (list[Tensor]): Selects the output targets of
+                the saved model. By default, all return variables of original
+                Layer's forward method are kept as the output of the saved model.
+                If the provided ``output_spec`` list is not all output variables,
+                the saved model will be pruned according to the given
+                ``output_spec`` list.
+
+        Returns:
+            None
+        """
+
+        assert isinstance(model, paddle.nn.Layer), (
+            "The model must be the instance of paddle.nn.Layer."
+        )
+
+        # Convert and save dygraph quantized model
+        self._convert(model)
+
+        paddle.jit.save(layer=model, path=path, input_spec=input_spec, **config)
+
+        # Load inference program
+        is_dynamic_mode = False
+        if paddle.in_dynamic_mode():
+            is_dynamic_mode = True
+            paddle.enable_static()
+
+        place = paddle.CPUPlace()
+        scope = paddle.static.global_scope()
+        exe = paddle.static.Executor(place)
+
+        dirname = os.path.dirname(path)
+        basename = os.path.basename(path)
+        model_filename = basename + INFER_MODEL_SUFFIX
+        params_filename = basename + INFER_PARAMS_SUFFIX
+
+        [
+            infer_program,
+            feed_target_names,
+            fetch_targets,
+        ] = paddle.static.load_inference_model(
+            path_prefix=dirname,
+            executor=exe,
+            model_filename=model_filename,
+            params_filename=params_filename,
+        )
+
+        # Process inference program
+        self._clean_up(infer_program)
+        self._gather_input_thresholds(infer_program, scope)
+        self._remove_scale_op(infer_program)
+
+        # Save final program
+        model_name = None
+        if model_filename is None:
+            model_name = "model"
+        elif model_filename.endswith(".pdmodel"):
+            model_name = model_filename.rsplit(".", 1)[0]
+        else:
+            model_name = model_filename
+        path_prefix = os.path.join(dirname, model_name)
+        feed_vars = [
+            infer_program.global_block().var(name) for name in feed_target_names
+        ]
+        paddle.static.save_inference_model(
+            path_prefix,
+            feed_vars,
+            fetch_targets,
+            executor=exe,
+            program=infer_program.clone(),
+        )
+
+        if is_dynamic_mode:
+            paddle.disable_static()
+
+    def _convert(self, model):
+        """
+        Convert the quantized model.
+
+        Args:
+            model(paddle.nn.Layer): The quantized model.
+            inplace(bool): Whether apply conversion to the input model.
+                           Default: False.
+        Returns:
+            None
+        """
+
+        for name, sub_layer in model.named_sublayers():
+            if self._is_quant_layer(sub_layer):
+                sub_layer._quant_config.quant_hook_handle.remove()
+
+        self._cal_thresholds(model)
+
+        for name, sub_layer in model.named_sublayers():
+            if self._is_quant_layer(sub_layer):
+                self._save_output_thresholds(sub_layer, sub_layer._quant_config)
+
+        self._wrap_simulated_layers(model)
+
+    def _cal_thresholds(self, model):
+        """
+        Calculate the thresholds of inputs and outputs.
+
+        Args:
+            model(paddle.nn.Layer): The quantized model.
+        Returns:
+            None
+        """
+        assert isinstance(model, paddle.nn.Layer), (
+            "The input model must be the instance of paddle.nn.Layer."
+        )
+
+        total_num = 0
+        cur_num = 0
+        for name, sub_layer in model.named_sublayers():
+            if self._is_quant_layer(sub_layer):
+                total_num += 1
+
+        for name, sub_layer in model.named_sublayers():
+            if self._is_quant_layer(sub_layer):
+                cur_num += 1
+                if cur_num % 5 == 0:
+                    _logger.info(f"Process the {cur_num} / {total_num} layer")
+
+                quant_config = sub_layer._quant_config
+
+                if quant_config.enable_in_act_quantizer:
+                    quant_config.in_act_quantizer.cal_thresholds()
+                quant_config.out_act_quantizer.cal_thresholds()
+
+                if PTQRegistry.is_simulated_quant_layer(sub_layer):
+                    weights = (sub_layer.weight,)
+                    quant_config.wt_quantizer.sample_data(sub_layer, weights)
+                    quant_config.wt_quantizer.cal_thresholds()
+
+    def _save_output_thresholds(self, sub_layer, quant_config):
+        """
+        Save the output thresholds to the layer.
+
+        Args:
+            sub_layer(paddle.nn.Layer): The quantized layer.
+            quant_config(PTQConfig): the quant config for the layer.
+        Returns:
+            None
+        """
+        assert isinstance(sub_layer, paddle.nn.Layer), (
+            "The input model must be the instance of paddle.nn.Layer."
+        )
+
+        layer_info = PTQRegistry.layer_info(sub_layer)
+
+        output_names = layer_info.output_names
+        output_thresholds = quant_config.out_act_quantizer.thresholds
+        assert len(output_names) == 1
+        if len(output_thresholds) == 1:
+            save_name = output_names[0] + str(0) + "_threshold"
+            sub_layer._set_op_attrs({save_name: output_thresholds[0]})
+            sub_layer._set_op_attrs({"out_threshold": output_thresholds[0]})
+        else:
+            _logger.warning(
+                f"output_thresholds shape of {output_names[0]} need to be 1, but received {len(output_thresholds)}"
+            )
+
+    def _wrap_simulated_layers(self, model):
+        """
+        Replace conv2d and linear with the quantized layers, and save
+        thresholds into the fake layers.
+        Args:
+            model(paddle.nn.Layer): The model to be quantized.
+        Returns:
+            None
+        """
+        assert isinstance(model, paddle.nn.Layer), (
+            "The input model must be the instance of paddle.nn.Layer."
+        )
+
+        for name, sub_layer in model.named_sublayers():
+            if self._is_quant_layer(
+                sub_layer
+            ) and PTQRegistry.is_simulated_quant_layer(sub_layer):
+                quant_config = sub_layer._quant_config
+                assert quant_config.enable_in_act_quantizer is True
+                wt_quantizer = quant_config.wt_quantizer
+                in_act_quantizer = quant_config.in_act_quantizer
+
+                # create layer
+                quant_layer_name = None
+                for key, value in utils.layer_name_map.items():
+                    if isinstance(sub_layer, value):
+                        quant_layer_name = 'Quantized' + key
+                        break
+                assert quant_layer_name is not None
+
+                if isinstance(wt_quantizer, ptq_quantizer.AbsmaxQuantizer):
+                    weight_quantize_type = "abs_max"
+                else:
+                    weight_quantize_type = "channel_wise_abs_max"
+                kwargs = {
+                    "weight_quantize_type": weight_quantize_type,
+                    "activation_quantize_type": "moving_average_abs_max",
+                    "weight_bits": wt_quantizer.quant_bits,
+                    "activation_bits": in_act_quantizer.quant_bits,
+                }
+
+                quant_layer = quant_layers.__dict__[quant_layer_name](
+                    sub_layer, **kwargs
+                )
+
+                # save the input thresholds
+                assert hasattr(quant_layer, "_fake_quant_input")
+                assert hasattr(quant_layer._fake_quant_input, "_scale")
+                if len(in_act_quantizer.thresholds) == 1:
+                    input_threshold = np.array(
+                        [in_act_quantizer.thresholds[0]], dtype=np.float32
+                    )
+                    quant_layer._fake_quant_input._scale.set_value(
+                        input_threshold
+                    )
+
+                assert hasattr(quant_layer, "_fake_quant_weight")
+                assert hasattr(quant_layer._fake_quant_weight, "_scale")
+                assert len(wt_quantizer.thresholds) == 1
+                weight_threshold = wt_quantizer.thresholds[0]
+                if isinstance(weight_threshold, list):
+                    weight_threshold = np.array(
+                        weight_threshold, dtype=np.float32
+                    )
+                else:
+                    weight_threshold = np.array(
+                        [weight_threshold], dtype=np.float32
+                    )
+                quant_layer._fake_quant_weight._scale.set_value(
+                    weight_threshold
+                )
+
+                # save the output thresholds
+                self._save_output_thresholds(quant_layer, quant_config)
+
+                # replace the layer
+                parent_layer, sub_name = utils.find_parent_layer_and_sub_name(
+                    model, name
+                )
+                setattr(parent_layer, sub_name, quant_layer)
+
+    def _gather_input_thresholds(self, program, scope):
+        """
+        Get and save input thresholds from the front ops.
+
+        Args:
+            program(Program): the input infer program.
+            scope(Scope): the corresponding scope for the program.
+        Returns:
+            None
+        """
+        for op in utils.program_all_ops(program):
+            for in_var_name in _get_op_input_var_names(op):
+                previous_op = utils.find_previous_op(op.block, in_var_name)
+                if previous_op is None:
+                    continue
+
+                if (
+                    "quantize_dequantize" in previous_op.type
+                    or previous_op.type == "moving_average_abs_max_scale"
+                ):
+                    attr_name = previous_op.output('OutScale')[0]
+                    in_threshold = utils.load_variable_data(scope, attr_name)
+                    in_threshold = utils.fp_numpy_to_naive(in_threshold)
+                    argname, index = _get_input_name_index(op, in_var_name)
+                    op._set_attr(
+                        argname + str(index) + "_threshold", in_threshold
+                    )
+                    op._set_attr("with_quant_attr", True)
+                else:
+                    for out_var_name in _get_op_output_var_names(previous_op):
+                        if out_var_name != in_var_name:
+                            continue
+                        argname, index = _get_output_name_index(
+                            previous_op, out_var_name
+                        )
+                        attr_name = argname + str(index) + "_threshold"
+                        if not previous_op.has_attr(attr_name):
+                            continue
+                        threshold = previous_op.attr(attr_name)
+
+                        argname, index = _get_input_name_index(op, in_var_name)
+                        attr_name = argname + str(index) + "_threshold"
+                        op._set_attr(attr_name, threshold)
+                        op._set_attr("with_quant_attr", True)
+
+    def _clean_up(self, program):
+        """
+        Remove useless thresholds which are added in jit.save.
+
+        Args:
+            program(Program): the input infer program.
+        Returns:
+            None
+        """
+
+        def _helper(op, next_op, old_attr_name, new_attr_name):
+            if (
+                op.has_attr(old_attr_name)
+                and next_op.has_attr(old_attr_name)
+                and op.attr(old_attr_name) == next_op.attr(old_attr_name)
+            ):
+                threshold = op.attr(old_attr_name)
+                op._remove_attr(old_attr_name)
+                next_op._remove_attr(old_attr_name)
+                next_op._set_attr(new_attr_name, threshold)
+                next_op._set_attr("with_quant_attr", True)
+
+        for op in utils.program_all_ops(program):
+            if "quantize_dequantize" in op.type:
+                # remove the thresholds in fake ops
+                for attr_name in op.attr_names:
+                    if "_threshold" in attr_name:
+                        op._remove_attr(attr_name)
+            elif op.type in ["conv2d", "matmul"]:
+                # change the thresholds in conv2d/matmul + eleadd
+                arg_name = "Output" if op.type == "conv2d" else "Out"
+                out_var_name = op.output(arg_name)[0]
+                next_ops = utils.find_next_ops(op.block, out_var_name)
+                if len(next_ops) > 1 or next_ops[0].type != "elementwise_add":
+                    continue
+                next_op = next_ops[0]
+
+                argname, index = _get_output_name_index(op, out_var_name)
+                old_attr_name = argname + str(index) + "_threshold"
+
+                argname, index = _get_output_name_index(
+                    next_op, next_op.output("Out")[0]
+                )
+                new_attr_name = argname + str(index) + "_threshold"
+
+                _helper(op, next_op, old_attr_name, new_attr_name)
+                _helper(op, next_op, "out_threshold", "out_threshold")
+
+    def _remove_scale_op(self, program):
+        """
+        Remove the moving_average_abs_max_scale op.
+        """
+        for op in utils.program_all_ops(program):
+            if op.type == "moving_average_abs_max_scale":
+                in_var_name = op.input("X")[0]
+                out_var_name = op.output("Out")[0]
+                next_ops = utils.find_next_ops(op.block, out_var_name)
+                for next_op in next_ops:
+                    next_op._rename_input(out_var_name, in_var_name)
+
+    @staticmethod
+    def _is_skip_layer(layer):
+        return hasattr(layer, "skip_quant") and layer.skip_quant is True
+
+    @staticmethod
+    def _is_quant_layer(layer):
+        return hasattr(layer, "_quant_config")

python/paddle/quantization/imperative/ptq_config.py

@@ -0,0 +1,55 @@
+#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+
+from .ptq_quantizer import (
+    SUPPORT_ACT_QUANTIZERS,
+    SUPPORT_WT_QUANTIZERS,
+    KLQuantizer,
+    PerChannelAbsmaxQuantizer,
+)
+
+
+class PTQConfig:
+    """
+    The PTQ config shows how to quantize the inputs and outputs.
+    """
+
+    def __init__(self, activation_quantizer, weight_quantizer):
+        """
+        Constructor.
+
+        Args:
+            activation_quantizer(BaseQuantizer): The activation quantizer.
+                It should be the instance of BaseQuantizer.
+            weight_quantizer(BaseQuantizer): The weight quantizer.
+                It should be the instance of BaseQuantizer.
+        """
+        super().__init__()
+        assert isinstance(activation_quantizer, tuple(SUPPORT_ACT_QUANTIZERS))
+        assert isinstance(weight_quantizer, tuple(SUPPORT_WT_QUANTIZERS))
+
+        self.in_act_quantizer = copy.deepcopy(activation_quantizer)
+        self.out_act_quantizer = copy.deepcopy(activation_quantizer)
+        self.wt_quantizer = copy.deepcopy(weight_quantizer)
+
+        self.quant_hook_handle = None
+
+        # In order to wrap simulated layers, use in_act_quantizer
+        # to calculate the input thresholds for conv2d, linear and etc.
+        self.enable_in_act_quantizer = False
+
+
+default_ptq_config = PTQConfig(KLQuantizer(), PerChannelAbsmaxQuantizer())

python/paddle/quantization/imperative/ptq_hooks.py

@@ -0,0 +1,27 @@
+#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+def quant_forward_post_hook(layer, inputs, outputs):
+    """
+    The forward_post_hook for PTQ.
+    """
+    assert hasattr(layer, '_quant_config'), (
+        "The layer should have _quant_config attr"
+    )
+
+    qc = layer._quant_config
+    if qc.enable_in_act_quantizer:
+        qc.in_act_quantizer.sample_data(layer, inputs)
+    qc.out_act_quantizer.sample_data(layer, (outputs,))

python/paddle/quantization/imperative/ptq_quantizer.py

@@ -0,0 +1,264 @@
+#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import abc
+import math
+
+import numpy as np
+
+import paddle
+
+from ...static.quantization.cal_kl_threshold import cal_kl_threshold
+from . import utils
+
+
+def abs_max_value(tensor):
+    return float(paddle.max(paddle.abs(tensor)))
+
+
+def merge_max_value(old, new):
+    """
+    Merge the max element one by one in two lists.
+    """
+    assert isinstance(old, list) and isinstance(new, list)
+    if old != []:
+        assert len(old) == len(new)
+        for i in range(len(old)):
+            assert type(old[i]) == type(new[i])
+            if isinstance(old[i], list):
+                new[i] = merge_max_value(old[i], new[i])
+            else:
+                new[i] = max(new[i], old[i])
+    return new
+
+
+def combine_abs_max_and_hist(
+    tensor, origin_max, origin_hist, bins, upsample_bins
+):
+    """ """
+
+    new_max = abs_max_value(tensor)
+
+    if new_max == 0.0:
+        return origin_max, origin_hist
+    elif origin_max == 0.0:
+        new_hist, _ = np.histogram(
+            paddle.abs(tensor).numpy(False), range=(0, new_max), bins=bins
+        )
+        new_hist = new_hist.astype(np.float32)
+        return new_max, new_hist
+    elif new_max <= origin_max:
+        new_hist, _ = np.histogram(
+            paddle.abs(tensor).numpy(False), range=(0, origin_max), bins=bins
+        )
+        new_hist = new_hist.astype(np.float32)
+        new_hist += origin_hist
+        return origin_max, new_hist
+    else:
+        # bin_width = origin_max / (bins * upsample_bins)
+        #           = new_max / (bins * downsample_bins)
+        bin_width = origin_max / (bins * upsample_bins)
+        downsample_bins = int(math.ceil(new_max / (bins * bin_width)))
+        new_max = bins * bin_width * downsample_bins
+
+        upsampled_hist = np.repeat(origin_hist, upsample_bins)
+        expanded_hist = np.zeros((bins * downsample_bins), dtype=np.float32)
+        expanded_hist[0 : bins * upsample_bins] = upsampled_hist
+        cumsumed_hist = np.cumsum(expanded_hist, dtype=np.float64)[
+            downsample_bins - 1 :: downsample_bins
+        ]
+        shift_cumsumed_hist = np.zeros((bins), dtype=np.float64)
+        shift_cumsumed_hist[1:] = cumsumed_hist[0:-1]
+        sampled_hist = (cumsumed_hist - shift_cumsumed_hist) / upsample_bins
+        sampled_hist = sampled_hist.astype(np.float32)
+
+        new_hist, _ = np.histogram(
+            paddle.abs(tensor).numpy(False), range=(0, new_max), bins=bins
+        )
+        new_hist = new_hist.astype(np.float32)
+        new_hist += sampled_hist
+
+        return new_max, new_hist
+
+
+class BaseQuantizer(metaclass=abc.ABCMeta):
+    """
+    Base quantizer for activation and weight.
+    """
+
+    def __init__(self, quant_bits=8):
+        super().__init__()
+        assert isinstance(quant_bits, int)
+        assert quant_bits > 0 and quant_bits <= 16
+
+        self.quant_bits = quant_bits
+
+        self.abs_max_vals = []
+        self.thresholds = []
+
+    @abc.abstractmethod
+    def sample_data(self, layer, tensors):
+        pass
+
+    @abc.abstractmethod
+    def cal_thresholds(self):
+        pass
+
+
+class AbsmaxQuantizer(BaseQuantizer):
+    """
+    Per-tensor abs max quantizer.
+    """
+
+    def __init__(self, quant_bits=8):
+        super().__init__(quant_bits)
+
+    def sample_data(self, layer, tensors):
+        assert isinstance(tensors, tuple)
+
+        abs_max_vals = [abs_max_value(t) for t in tensors]
+        self.abs_max_vals = merge_max_value(self.abs_max_vals, abs_max_vals)
+
+    def cal_thresholds(self):
+        self.thresholds = self.abs_max_vals
+
+
+class PerChannelAbsmaxQuantizer(BaseQuantizer):
+    """
+    Per channel abs max quantizer.
+    """
+
+    def __init__(self, quant_bits=8):
+        super().__init__(quant_bits)
+
+    def sample_data(self, layer, tensors):
+        assert isinstance(layer, paddle.nn.Layer)
+        assert isinstance(tensors, tuple)
+
+        abs_max_vals_list = []
+        for idx, tensor in enumerate(tensors):
+            if isinstance(layer, tuple(utils.spec_channel_axis_layers)):
+                abs_max_vals = [
+                    abs_max_value(tensor[:, i]) for i in range(tensor.shape[1])
+                ]
+                abs_max_vals_list.append(abs_max_vals)
+            else:
+                abs_max_vals = [
+                    abs_max_value(tensor[i]) for i in range(tensor.shape[0])
+                ]
+                abs_max_vals_list.append(abs_max_vals)
+
+        self.abs_max_vals = merge_max_value(
+            self.abs_max_vals, abs_max_vals_list
+        )
+
+    def cal_thresholds(self):
+        self.thresholds = self.abs_max_vals
+
+
+class BaseHistQuantizer(BaseQuantizer, metaclass=abc.ABCMeta):
+    """ """
+
+    def __init__(self, quant_bits=8, bins=1024, upsample_bins=64):
+        super().__init__(quant_bits)
+        self.bins = bins
+        self.upsample_bins = upsample_bins
+
+        self.hists = []
+
+    def sample_data(self, layer, tensors):
+        assert isinstance(tensors, tuple)
+
+        if self.abs_max_vals == []:
+            abs_max_vals = [abs_max_value(t) for t in tensors]
+            self.abs_max_vals = abs_max_vals
+
+            for idx, tensor in enumerate(tensors):
+                if abs_max_vals[idx] == 0.0:
+                    self.hists.append(None)
+                else:
+                    hist, _ = np.histogram(
+                        paddle.abs(tensor).numpy(False),
+                        range=(0.0, abs_max_vals[idx]),
+                        bins=self.bins,
+                    )
+                    hist = hist.astype(np.float32)
+                    self.hists.append(hist)
+        else:
+            assert len(self.abs_max_vals) == len(tensors)
+            assert len(self.hists) == len(tensors)
+
+            for idx, tensor in enumerate(tensors):
+                new_abs_max, new_hist = combine_abs_max_and_hist(
+                    tensor,
+                    self.abs_max_vals[idx],
+                    self.hists[idx],
+                    self.bins,
+                    self.upsample_bins,
+                )
+                self.abs_max_vals[idx] = new_abs_max
+                self.hists[idx] = new_hist
+
+    @abc.abstractmethod
+    def cal_thresholds(self):
+        pass
+
+
+class HistQuantizer(BaseHistQuantizer):
+    """ """
+
+    def __init__(
+        self, quant_bits=8, bins=1024, upsample_bins=64, hist_percent=0.99999
+    ):
+        super().__init__(quant_bits, bins, upsample_bins)
+        self.hist_percent = hist_percent
+
+    def cal_thresholds(self):
+        def _helper(abs_max, hist, percent):
+            assert hist.ndim == 1 and percent < 1.0
+            hist = hist / np.sum(hist, dtype=np.float64)
+            cumsumed_hist = np.cumsum(hist)
+            index = np.argwhere(cumsumed_hist >= percent)[0]
+            return float((index - 0.5) * (abs_max / hist.shape[0]))
+
+        for idx in range(len(self.hists)):
+            if self.hists[idx] is None:
+                self.thresholds.append(self.abs_max_vals[idx])
+            else:
+                threshold = _helper(
+                    self.abs_max_vals[idx], self.hists[idx], self.hist_percent
+                )
+                self.thresholds.append(threshold)
+
+
+class KLQuantizer(BaseHistQuantizer):
+    """ """
+
+    def __init__(self, quant_bits=8, bins=1024, upsample_bins=64):
+        super().__init__(quant_bits, bins, upsample_bins)
+
+    def cal_thresholds(self):
+        for idx in range(len(self.hists)):
+            if self.hists[idx] is None:
+                self.thresholds.append(self.abs_max_vals[idx])
+            else:
+                hist = self.hists[idx]
+                abs_max_val = self.abs_max_vals[idx]
+                bin_width = abs_max_val / hist.shape[0]
+                threshold = cal_kl_threshold(hist, bin_width, self.quant_bits)
+                self.thresholds.append(threshold)
+
+
+SUPPORT_ACT_QUANTIZERS = [AbsmaxQuantizer, HistQuantizer, KLQuantizer]
+SUPPORT_WT_QUANTIZERS = [AbsmaxQuantizer, PerChannelAbsmaxQuantizer]

python/paddle/quantization/imperative/ptq_registry.py

@@ -0,0 +1,143 @@
+#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import paddle
+
+
+class LayerInfo:
+    """
+    Store the arg names of the inputs and outputs.
+    """
+
+    def __init__(self, layer, input_names, weight_names, output_names):
+        super().__init__()
+        self.layer = layer
+        self.input_names = input_names
+        self.weight_names = weight_names
+        self.output_names = output_names
+
+
+PTQ_LAYERS_INFO = [
+    LayerInfo(paddle.nn.Conv2D, ['Input'], ['Filter'], ['Output']),
+    LayerInfo(paddle.nn.Linear, ['X'], ['Y'], ['Out']),
+    LayerInfo(paddle.nn.BatchNorm2D, ['X'], [], ['Y']),
+    LayerInfo(paddle.nn.AdaptiveMaxPool2D, ['X'], [], ['Out']),
+    LayerInfo(paddle.nn.AdaptiveAvgPool2D, ['X'], [], ['Out']),
+    LayerInfo(paddle.nn.AvgPool2D, ['X'], [], ['Out']),
+    LayerInfo(paddle.nn.MaxPool2D, ['X'], [], ['Out']),
+    LayerInfo(paddle.nn.ReLU, ['X'], [], ['Out']),
+    LayerInfo(paddle.nn.ReLU6, ['X'], [], ['Out']),
+    LayerInfo(paddle.nn.Hardswish, ['X'], [], ['Out']),
+    LayerInfo(paddle.nn.Swish, ['X'], [], ['Out']),
+    LayerInfo(paddle.nn.Sigmoid, ['X'], [], ['Out']),
+    LayerInfo(paddle.nn.Softmax, ['X'], [], ['Out']),
+    LayerInfo(paddle.nn.Tanh, ['X'], [], ['Out']),
+    LayerInfo(paddle.nn.quant.add, ['X', 'Y'], [], ['Out']),
+]
+
+QUANT_LAYERS_INFO = [
+    LayerInfo(
+        paddle.nn.quant.quant_layers.QuantizedConv2D,
+        ['Input'],
+        ['Filter'],
+        ['Output'],
+    ),
+    LayerInfo(
+        paddle.nn.quant.quant_layers.QuantizedLinear, ['X'], ['Y'], ['Out']
+    ),
+]
+
+SIMULATED_LAYERS = [paddle.nn.Conv2D, paddle.nn.Linear]
+
+
+class PTQRegistry:
+    """
+    Register the supported layers for PTQ and provide layers info.
+    """
+
+    supported_layers_map = {}
+    registered_layers_map = {}
+    is_inited = False
+
+    def __init__(self):
+        super().__init__()
+
+    @classmethod
+    def _init(cls):
+        if not cls.is_inited:
+            for layer_info in PTQ_LAYERS_INFO:
+                cls.supported_layers_map[layer_info.layer] = layer_info
+
+            all_layers_info = PTQ_LAYERS_INFO + QUANT_LAYERS_INFO
+            for layer_info in all_layers_info:
+                cls.registered_layers_map[layer_info.layer] = layer_info
+        cls.is_inited = True
+
+    @classmethod
+    def is_supported_layer(cls, layer):
+        """
+        Analyze whether the layer supports quantization.
+        Args:
+            layer(Layer): The input layer can be a python class or an instance.
+        Returns:
+            flag(bool): Whether the layer is supported.
+        """
+        cls._init()
+        return layer in cls.supported_layers_map or isinstance(
+            layer, tuple(cls.supported_layers_map.keys())
+        )
+
+    @classmethod
+    def is_registered_layer(cls, layer):
+        """
+        Analyze whether the layer is register layer_info.
+        Args:
+            layer(Layer): The input layer can be a python class or an instance.
+        Returns:
+            flag(bool): Whether the layer is register layer_info.
+        """
+        cls._init()
+        return layer in cls.registered_layers_map or isinstance(
+            layer, tuple(cls.registered_layers_map.keys())
+        )
+
+    @classmethod
+    def is_simulated_quant_layer(cls, layer):
+        """
+        Analyze whether the layer is simulated quant layer.
+        Args:
+            layer(Layer): The input layer can be a python class or an instance.
+        Returns:
+            flag(bool): Whether the layer is supported.
+        """
+        return layer in SIMULATED_LAYERS or isinstance(
+            layer, tuple(SIMULATED_LAYERS)
+        )
+
+    @classmethod
+    def layer_info(cls, layer):
+        """
+        Get the information for the layer.
+        Args:
+            layer(Layer): The input layer can be a python class or an instance.
+        Returns:
+            layer_info(LayerInfo): The layer info of the input layer.
+        """
+        assert cls.is_registered_layer(layer), (
+            "The input layer is not register."
+        )
+
+        for layer_key, layer_info in cls.registered_layers_map.items():
+            if layer == layer_key or isinstance(layer, layer_key):
+                return layer_info

python/paddle/quantization/imperative/qat.py

@@ -0,0 +1,765 @@
+#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+import paddle
+from paddle.base.framework import IrGraph
+from paddle.framework import core
+from paddle.nn.quant import quant_layers
+
+from ...static.quantization.quantization_pass import (
+    QuantWeightPass,
+    ReplaceFakeQuantDequantPass,
+)
+from ...static.quantization.utils import (
+    _get_input_name_index,
+    _get_op_input_var_names,
+    _get_output_name_index,
+    move_persistable_var_to_global_block,
+)
+from . import fuse_utils, utils
+
+INFER_MODEL_SUFFIX = ".pdmodel"
+INFER_PARAMS_SUFFIX = ".pdiparams"
+
+
+def lazy_import_fleet(layer_name_map, fake_quant_input_layers):
+    from paddle.distributed import fleet
+
+    layer_name_map['ColumnParallelLinear'] = (
+        fleet.meta_parallel.parallel_layers.mp_layers.ColumnParallelLinear
+    )
+    layer_name_map['RowParallelLinear'] = (
+        fleet.meta_parallel.parallel_layers.mp_layers.RowParallelLinear
+    )
+    fake_quant_input_layers.append(fleet.meta_parallel.RowParallelLinear)
+    fake_quant_input_layers.append(fleet.meta_parallel.ColumnParallelLinear)
+    return layer_name_map, fake_quant_input_layers
+
+
+class ImperativeQuantAware:
+    """
+    Applying quantization aware training (QAT) to the dygraph model.
+    """
+
+    def __init__(
+        self,
+        quantizable_layer_type=[
+            'Conv2D',
+            'Linear',
+            'Conv2DTranspose',
+            'ColumnParallelLinear',
+            'RowParallelLinear',
+        ],
+        weight_quantize_type='abs_max',
+        activation_quantize_type='moving_average_abs_max',
+        weight_bits=8,
+        activation_bits=8,
+        moving_rate=0.9,
+        fuse_conv_bn=False,
+        weight_preprocess_layer=None,
+        act_preprocess_layer=None,
+        weight_quantize_layer=None,
+        act_quantize_layer=None,
+        onnx_format=False,
+    ):
+        """
+        The constructor for ImperativeQuantAware.
+
+        Args:
+            quantizable_layer_type(list[str | layer]): List the type of
+                layers that will be quantized. Default is ['Conv2D', 'Linear'].
+            weight_quantize_type(str): quantization type for weights,
+                which supports 'abs_max' and 'channel_wise_abs_max'.
+            activation_quantize_type(str): quantization type for activations,
+                which supports 'abs_max' and 'moving_average_abs_max' now.
+                If using 'abs_max' mode, the quantization scale will be
+                calculated dynamically each step in both training and testing
+                period. If using 'moving_average_abs_max', the static
+                quantization scale will be calculated during training and
+                used in inference.
+            weight_bits(int): quantization bit number for weights, whereas
+                the bias is not quantized.
+            activation_bits(int): quantization bit number for activations.
+            moving_rate(float): the parameter for 'moving_average_abs_max'
+                quantization.
+            fuse_conv_bn(bool): Whether to fuse conv and bn, default is False.
+            weight_preprocess_layer(paddle.nn.Layer, optional): A paddle
+                Layer that defines how to preprocess weight before quantization.
+                Using this can quickly test if user's preprocess method works
+                or not. The input is non-quantized weight and function returns
+                processed weight to be quantized.
+                If None, the weight will be quantized directly.
+                Default is None.
+            act_preprocess_layer(paddle.nn.Layer, optional): A paddle Layer
+                that defines how to preprocess activation before quantization.
+                Using this can quickly test if user's preprocess method works
+                or not. The input is non-quantized activation and function returns
+                processed activation to be quantized.
+                If None, the activation will be quantized directly.
+                Default is None.
+            weight_quantize_layer(paddle.nn.Layer, optional): A paddle Layer that
+                defines how to quantize weight.
+                Using this can quickly test if user's quantization method works or not.
+                In this layer, user should both define quantization method and
+                dequantization method, that is, the function's input is non-quantized
+                weight and returns dequantized weight.
+                If None, will use quantization op defined by 'weight_quantize_type'.
+                Default is None.
+            act_quantize_layer(paddle.nn.Layer, optional): A paddle Layer that defines
+                how to quantize activation.
+                Using this can quickly test if user's quantization method works or not.
+                In this layer, user should both define quantization method and
+                dequantization method, that is, the function's input is non-quantized
+                activation and returns dequantized activation.
+                If None, will use quantization op defined by 'activation_quantize_type'.
+                Default is None.
+            onnx_format (bool, optional): Whether to export the quantized model
+                with format of ONNX. Default is False.
+
+        Note:
+            If user sets attribute 'skip_quant' to a Layer that support dynamic
+            quantization and sets it to true, the layer would not be quantized
+            during training. If this attribute is not sets or the attribute is
+            false, the Layer would be quantized in training.
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle
+                >>> from paddle.static.quantization import (
+                ...     ImperativeQuantAware,
+                ... )
+                >>> from paddle.vision.models import (
+                ...     resnet,
+                ... )
+
+                >>> model = resnet.resnet50(pretrained=True)
+
+                >>> imperative_qat = ImperativeQuantAware(
+                ...     weight_quantize_type='abs_max',
+                ...     activation_quantize_type='moving_average_abs_max',
+                ... )
+
+                >>> # Add the fake quant logical.
+                >>> # The original model will be rewrite.
+                >>> # The outscale of outputs in supported layers would be calculated.
+                >>> imperative_qat.quantize(model)
+
+                >>> # Fine-tune the quantized model
+                >>> # ...
+
+                >>> # Save quant model for the inference.
+                >>> imperative_qat.save_quantized_model(
+                ...     layer=model,
+                ...     model_path="./resnet50_qat",
+                ...     input_spec=[
+                ...         paddle.static.InputSpec(shape=[None, 3, 224, 224], dtype='float32'),
+                ...     ],
+                ... )
+
+            .. code-block:: pycon
+
+                >>> import paddle
+                >>> from paddle.static.quantization import (
+                ...     ImperativeQuantAware,
+                ... )
+
+                >>> class ImperativeModel(paddle.nn.Layer):
+                ...     def __init__(self):
+                ...         super().__init__()
+                ...         # self.linear_0 would skip the quantization.
+                ...         self.linear_0 = paddle.nn.Linear(784, 400)
+                ...         self.linear_0.skip_quant = True
+
+                ...         # self.linear_1 would not skip the quantization.
+                ...         self.linear_1 = paddle.nn.Linear(400, 10)
+                ...         self.linear_1.skip_quant = False
+
+                ...     def forward(self, inputs):
+                ...         x = self.linear_0(inputs)
+                ...         x = self.linear_1(inputs)
+                ...         return x
+
+                >>> model = ImperativeModel()
+                >>> imperative_qat = ImperativeQuantAware(
+                ...     weight_quantize_type='abs_max',
+                ...     activation_quantize_type='moving_average_abs_max',
+                ... )
+
+                >>> # Add the fake quant logical.
+                >>> # The original model will be rewrite.
+                >>> #
+                >>> # There is only one Layer(self.linear1) would be added the
+                >>> # fake quant logical.
+                >>> imperative_qat.quantize(model)
+
+                >>> # Fine-tune the quantized model
+                >>> # ...
+
+                >>> # Save quant model for the inference.
+                >>> imperative_qat.save_quantized_model(
+                ...     layer=model,
+                ...     model_path="./imperative_model_qat",
+                ... )
+        """
+        super().__init__()
+        self.fuse_conv_bn = fuse_conv_bn
+
+        kwargs = {
+            "quantizable_layer_type": quantizable_layer_type,
+            "weight_quantize_type": weight_quantize_type,
+            "activation_quantize_type": activation_quantize_type,
+            "weight_bits": weight_bits,
+            "activation_bits": activation_bits,
+            "moving_rate": moving_rate,
+            "weight_preprocess_layer": weight_preprocess_layer,
+            "act_preprocess_layer": act_preprocess_layer,
+            "weight_quantize_layer": weight_quantize_layer,
+            "act_quantize_layer": act_quantize_layer,
+        }
+
+        self._quantize_inputs = ImperativeQuantizeInputs(**kwargs)
+
+        self._quantize_outputs = ImperativeQuantizeOutputs(
+            moving_rate, activation_bits, onnx_format
+        )
+
+    def quantize(self, model):
+        """
+        According to weights' and activations' quantization types,
+        the model will be added some fake quant ops, such as
+        fake_quantize_dequantize_moving_average_abs_max,
+        fake_quantize_dequantize_abs_max and so on. At the same time,
+        the out_scale value of outputs would be calculated.
+
+        Args:
+            model(paddle.nn.Layer): the model to be quantized.
+        Returns:
+            None
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle
+                >>> from paddle.static.quantization import (
+                ...     ImperativeQuantAware,
+                ... )
+
+                >>> class ImperativeModel(paddle.nn.Layer):
+                ...     def __init__(self):
+                ...         super().__init__()
+                ...         # self.linear_0 would skip the quantization.
+                ...         self.linear_0 = paddle.nn.Linear(784, 400)
+                ...         self.linear_0.skip_quant = True
+
+                ...         # self.linear_1 would not skip the quantization.
+                ...         self.linear_1 = paddle.nn.Linear(400, 10)
+                ...         self.linear_1.skip_quant = False
+
+                ...     def forward(self, inputs):
+                ...         x = self.linear_0(inputs)
+                ...         x = self.linear_1(inputs)
+                ...         return x
+
+                >>> model = ImperativeModel()
+                >>> imperative_qat = ImperativeQuantAware(
+                ...     weight_quantize_type='abs_max',
+                ...     activation_quantize_type='moving_average_abs_max',
+                ... )
+
+                >>> # Add the fake quant logical.
+                >>> # The original model will be rewrite.
+                >>> #
+                >>> # There is only one Layer(self.linear1) would be added the
+                >>> # fake quant logical.
+                >>> imperative_qat.quantize(model)
+        """
+        assert isinstance(model, paddle.nn.Layer), (
+            "The model must be the instance of paddle.nn.Layer."
+        )
+
+        if self.fuse_conv_bn:
+            fuse_utils.fuse_conv_bn(model)
+
+        self._quantize_inputs.apply(model)
+        self._quantize_outputs.apply(model)
+        return model
+
+    def save_quantized_model(self, layer, path, input_spec=None, **config):
+        with paddle.pir_utils.OldIrGuard():
+            self._quantize_outputs.save_quantized_model(
+                layer, path, input_spec, **config
+            )
+
+
+class ImperativeQuantizeInputs:
+    """
+    Based on the input params, add the quant_dequant computational
+    logic both for activation inputs and weight inputs.
+    """
+
+    def __init__(
+        self,
+        quantizable_layer_type=['Conv2D', 'Linear', 'Conv2DTranspose'],
+        weight_quantize_type='abs_max',
+        activation_quantize_type='moving_average_abs_max',
+        weight_bits=8,
+        activation_bits=8,
+        moving_rate=0.9,
+        weight_preprocess_layer=None,
+        act_preprocess_layer=None,
+        weight_quantize_layer=None,
+        act_quantize_layer=None,
+    ):
+        """
+        The constructor for ImperativeQuantizeInputs.
+
+        Please refer to the args of ImperativeQuantAware.
+        """
+        super().__init__()
+        self.layer_name_map, self.fake_quant_input_layers = lazy_import_fleet(
+            utils.layer_name_map, utils.fake_quant_input_layers
+        )
+
+        self._quantizable_layer_type = tuple(
+            (
+                self.layer_name_map[layer]
+                if layer in self.layer_name_map
+                else layer
+            )
+            for layer in quantizable_layer_type
+        )
+        for layer in self._quantizable_layer_type:
+            assert (
+                not isinstance(layer, str)
+                and layer in self.fake_quant_input_layers
+            ), f"{layer} is unsupported to be quantized."
+
+        quantize_type = {
+            'abs_max',
+            'moving_average_abs_max',
+            'channel_wise_abs_max',
+            'lsq_weight',
+            'channel_wise_lsq_weight',
+        }
+        act_quantize_type = {'moving_average_abs_max', 'lsq_act'}
+        assert (
+            weight_quantize_type != 'moving_average_abs_max'
+            and weight_quantize_type in quantize_type
+        ), (
+            f"Unsupported weight_quantize_type: {weight_quantize_type}. It can only "
+            "be abs_max or channel_wise_abs_max."
+        )
+        # TODO (jc): activation_quantize_type supports range_abs_max
+        assert activation_quantize_type in act_quantize_type, (
+            f"Unsupported activation_quantize_type: {activation_quantize_type}. It can "
+            "only be moving_average_abs_max or lsq_act now."
+        )
+
+        bits_check = lambda bits: (
+            isinstance(bits, int) and bits >= 0 and bits <= 16
+        )
+        assert bits_check(weight_bits), "weight_bits should be 1, 2,... or 16."
+        assert bits_check(activation_bits), (
+            "activation_bits should be 1, 2,... or 16."
+        )
+
+        layer_check = lambda method: (
+            method is None or issubclass(method, paddle.nn.Layer)
+        )
+        assert layer_check(weight_preprocess_layer), (
+            "weight_preprocess should be nn.Layer."
+        )
+        assert layer_check(act_preprocess_layer), (
+            "act_preprocess should be nn.Layer."
+        )
+        assert layer_check(weight_quantize_layer), (
+            "weight_quantize should be nn.Layer."
+        )
+        assert layer_check(act_quantize_layer), (
+            "act_quantize should be nn.Layer."
+        )
+
+        self._kwargs = {
+            "weight_quantize_type": weight_quantize_type,
+            "activation_quantize_type": activation_quantize_type,
+            "weight_bits": weight_bits,
+            "activation_bits": activation_bits,
+            "moving_rate": moving_rate,
+            "weight_pre_layer": weight_preprocess_layer,
+            "act_pre_layer": act_preprocess_layer,
+            "weight_quant_layer": weight_quantize_layer,
+            "act_quant_layer": act_quantize_layer,
+        }
+
+    def apply(self, model):
+        """
+        Quantize the weights and activations to calculate for specific
+        layers.
+
+        Args:
+            model(paddle.nn.Layer): The target model which would
+                calculate the input quantization scale.
+
+        Returns:
+            None
+        """
+
+        assert isinstance(model, paddle.nn.Layer), (
+            "The model must be the instance of paddle.nn.Layer."
+        )
+
+        for name, cur_layer in model.named_sublayers():
+            if not isinstance(cur_layer, self._quantizable_layer_type) or (
+                hasattr(cur_layer, "skip_quant")
+                and cur_layer.skip_quant is True
+            ):
+                continue
+
+            parent_layer, sub_name = utils.find_parent_layer_and_sub_name(
+                model, name
+            )
+
+            cur_quant_layer = self._get_input_quantized_layer(cur_layer)
+            setattr(parent_layer, sub_name, cur_quant_layer)
+
+    def _get_input_quantized_layer(self, layer):
+        quant_layer_name = None
+
+        for key, value in self.layer_name_map.items():
+            if isinstance(layer, value):
+                quant_layer_name = 'Quantized' + key
+                break
+        assert quant_layer_name is not None, (
+            f"The layer {layer.full_name()} is unsupported to be quantized."
+        )
+
+        return quant_layers.__dict__[quant_layer_name](layer, **self._kwargs)
+
+
+class ImperativeQuantizeOutputs:
+    """
+    Calculate the output scales for target layers.
+    """
+
+    def __init__(self, moving_rate=0.9, activation_bits=8, onnx_format=False):
+        """
+        The constructor for ImperativeQuantizeOutputs.
+
+        Args:
+            moving_rate(float): The decay coefficient of moving average.
+                                The default value is 0.9.
+            activation_bits(int, optional): quantization bit number for activation. Default is 8.
+        """
+        super().__init__()
+        self._moving_rate = moving_rate
+        self._activation_bits = activation_bits
+        self._onnx_format = onnx_format
+
+    def apply(self, model):
+        """
+        Insert the `moving_average_abs_max_scale` layers to calculate the
+        output scales for specific layers in the dygraph model.
+
+        Args:
+            model(paddle.nn.Layer): The target model which would be
+                calculate the output quantization scale.
+
+        Returns:
+            None
+        """
+        assert isinstance(model, paddle.nn.Layer), (
+            "The model must be the instance of paddle.nn.Layer."
+        )
+
+        for cur_name, cur_layer in model.named_sublayers():
+            if '_act_preprocess' in cur_name:
+                continue
+            if not self._is_target_layer(cur_layer):
+                continue
+
+            parent_layer, sub_name = utils.find_parent_layer_and_sub_name(
+                model, cur_name
+            )
+
+            reduce_type = None
+
+            if isinstance(cur_layer, tuple(utils.fake_quant_output_layers)):
+                cur_quant_layer = quant_layers.FakeQuantMAOutputScaleLayer(
+                    cur_layer, self._moving_rate, reduce_type=reduce_type
+                )
+            else:
+                cur_quant_layer = quant_layers.MAOutputScaleLayer(
+                    cur_layer, self._moving_rate, reduce_type=reduce_type
+                )
+
+            setattr(parent_layer, sub_name, cur_quant_layer)
+
+    def save_quantized_model(self, model, path, input_spec=None, **config):
+        """
+        Save the quantized model for the inference.
+
+        Args:
+            model (Layer): The model to be saved.
+            path (str): The path prefix to save model. The format is
+                ``dirname/file_prefix`` or ``file_prefix``.
+            input_spec (list[InputSpec|Tensor], optional): Describes the input
+                of the saved model's forward method, which can be described by
+                InputSpec or example Tensor. If None, all input variables of
+                the original Layer's forward method would be the inputs of
+                the saved model. Default None.
+            **config (dict, optional): Other save configuration options for
+                compatibility. We do not recommend using these configurations,
+                they may be removed in the future. If not necessary, DO NOT use
+                them. Default None.
+                The following options are currently supported:
+                (1) output_spec (list[Tensor]): Selects the output targets of
+                the saved model. By default, all return variables of original
+                Layer's forward method are kept as the output of the saved model.
+                If the provided ``output_spec`` list is not all output variables,
+                the saved model will be pruned according to the given
+                ``output_spec`` list.
+
+        Returns:
+            None
+        """
+        assert isinstance(model, paddle.nn.Layer), (
+            "The model must be the instance of paddle.nn.Layer."
+        )
+
+        if input_spec:
+            paddle.jit.to_static(model, input_spec=input_spec)
+        paddle.jit.save(layer=model, path=path, input_spec=input_spec, **config)
+
+        is_dynamic_mode = False
+        if paddle.in_dynamic_mode():
+            is_dynamic_mode = True
+            paddle.enable_static()
+
+        place = core.CPUPlace()
+        scope = paddle.static.global_scope()
+        exe = paddle.static.Executor(place)
+
+        dirname = os.path.dirname(path)
+        basename = os.path.basename(path)
+        model_filename = basename + INFER_MODEL_SUFFIX
+        params_filename = basename + INFER_PARAMS_SUFFIX
+
+        [
+            infer_program,
+            feed_target_names,
+            fetch_targets,
+        ] = paddle.static.load_inference_model(
+            dirname,
+            executor=exe,
+            model_filename=model_filename,
+            params_filename=params_filename,
+        )
+
+        if not self._onnx_format:
+            self._gather_scales(infer_program, scope, fetch_targets)
+
+            # Remove `moving_average_abs_max_scale` node in sub graphs.
+            graph = IrGraph(core.Graph(infer_program.desc), for_test=False)
+            for sub_graph in graph.all_sub_graphs():
+                for _op in sub_graph.all_op_nodes():
+                    if _op.name() == "moving_average_abs_max_scale":
+                        sub_graph.safe_remove_nodes(_op)
+                sub_graph.resolve_hazard()
+            infer_program = graph.to_program()
+
+            self._set_skip_quant_attr(infer_program)
+
+            clip_extra = False
+        else:
+            graph = IrGraph(core.Graph(infer_program.desc), for_test=False)
+            transform_pass = ReplaceFakeQuantDequantPass(
+                scope, place, quant_bits=self._activation_bits
+            )
+            for sub_graph in graph.all_sub_graphs():
+                sub_graph._for_test = True
+                transform_pass.apply(sub_graph)
+
+            quant_weight_pass = QuantWeightPass(scope, place)
+            for sub_graph in graph.all_sub_graphs():
+                sub_graph._for_test = True
+                quant_weight_pass.apply(sub_graph)
+
+            infer_program = graph.to_program()
+
+            clip_extra = True
+
+        move_persistable_var_to_global_block(infer_program)
+
+        model_name = None
+        if model_filename is None:
+            model_name = "model"
+        elif model_filename.endswith(".pdmodel"):
+            model_name = model_filename.rsplit(".", 1)[0]
+        else:
+            model_name = model_filename
+        path_prefix = os.path.join(dirname, model_name)
+        feed_vars = [
+            infer_program.global_block().var(name) for name in feed_target_names
+        ]
+        paddle.static.save_inference_model(
+            path_prefix,
+            feed_vars,
+            fetch_targets,
+            executor=exe,
+            program=infer_program.clone(),
+            clip_extra=clip_extra,
+        )
+
+        if is_dynamic_mode:
+            paddle.disable_static()
+
+    def _is_target_layer(self, layer):
+        """
+        Whether the layer needs to calculate output scales.
+        """
+        # exclude fake_quant ops in quant_layers file
+        if not isinstance(layer, paddle.nn.Layer):
+            return False
+
+        if self._onnx_format:
+            return (
+                True
+                if isinstance(layer, tuple(utils.fake_quant_wrap_layers))
+                else False
+            )
+
+        flag = False
+        if utils.is_leaf_layer(layer) and not isinstance(
+            layer, tuple(utils.fake_quant_leaf_layers)
+        ):
+            flag = True
+
+        if isinstance(layer, tuple(utils.fake_quant_wrap_layers)):
+            flag = True
+
+        if isinstance(layer, paddle.nn.quant.FloatFunctionalLayer):
+            flag = True
+
+        return flag
+
+    def _gather_scales(self, program, scope, fetch_targets):
+        """
+        Get all scales from fake ops, save them into the corresponding ops
+        and delete all moving_average_abs_max_scale ops.
+        """
+
+        def _gather_input_scale():
+            target_ops = []
+            skip_ops = [
+                *utils.fake_quantize_dequantize_op_types,
+                "moving_average_abs_max_scale",
+            ]
+            for block in program.blocks:
+                for op in block.ops:
+                    if op.type not in skip_ops:
+                        target_ops.append(op)
+
+            for op in target_ops:
+                for in_var_name in _get_op_input_var_names(op):
+                    previous_op = utils.find_previous_op(op.block, in_var_name)
+
+                    if previous_op is not None and (
+                        "quantize_dequantize" in previous_op.type
+                        or previous_op.type == "moving_average_abs_max_scale"
+                    ):
+                        scale_name = previous_op.output('OutScale')[0]
+                        in_scale = utils.load_variable_data(scope, scale_name)
+                        in_scale = utils.fp_numpy_to_naive(in_scale)
+                        argname, index = _get_input_name_index(op, in_var_name)
+                        op._set_attr(
+                            argname + str(index) + "_threshold", in_scale
+                        )
+                        op._set_attr("with_quant_attr", True)
+
+        def _gather_output_scale():
+            target_ops = []
+            for block in program.blocks:
+                for op in block.ops:
+                    if op.type == "moving_average_abs_max_scale":
+                        target_ops.append(op)
+
+            for op in target_ops:
+                in_var_name = op.input('X')[0]
+                out_var_name = op.output('Out')[0]
+                block = op.block
+                previous_op = utils.find_previous_op(block, in_var_name)
+                next_ops = utils.find_next_ops(block, out_var_name)
+
+                out_scale_name = op.output('OutScale')[0]
+                out_scale = utils.load_variable_data(scope, out_scale_name)
+                out_scale = utils.fp_numpy_to_naive(out_scale)
+
+                if previous_op.type != "feed":
+                    res = _get_output_name_index(previous_op, in_var_name)
+                    if res is not None:
+                        argname, index = res
+                        previous_op._set_attr(
+                            argname + str(index) + "_threshold", out_scale
+                        )
+                        previous_op._set_attr("out_threshold", out_scale)
+                        previous_op._set_attr("with_quant_attr", True)
+
+                for next_op in next_ops:
+                    next_op._rename_input(out_var_name, in_var_name)
+                    # If next_op is `fetch` and out_var_name in fetch_targets,
+                    # fetch_targets must update to in_var_name when rename input.
+                    for i in range(len(fetch_targets)):
+                        if fetch_targets[i].name == out_var_name:
+                            fetch_targets[i] = block.var(in_var_name)
+
+        _gather_input_scale()
+        _gather_output_scale()
+
+    def _set_skip_quant_attr(self, program):
+        """
+        Label the skip quantized ops.
+        """
+        for block in program.blocks:
+            for op in block.ops:
+                if self._is_skip_quant_op(block, op):
+                    op._set_attr("skip_quant", True)
+                    op._set_attr("with_quant_attr", True)
+
+    def _is_skip_quant_op(self, block, in_op):
+        """
+        The input op should be skipped quantization.
+        1. the type of input op should be conv2d, depthwise_conv2d or matmul
+        2. the previous ops of the input op are not fake_quantize_dequantize ops
+        """
+        target_op_types = [
+            "conv2d",
+            "depthwise_conv2d",
+            "matmul",
+            "conv2d_transpose",
+        ]
+        if in_op.type not in target_op_types:
+            return False
+
+        previous_ops = [
+            utils.find_previous_op(block, arg_name)
+            for arg_name in in_op.input_arg_names
+        ]
+        return any(
+            op is not None
+            and op.type not in utils.fake_quantize_dequantize_op_types
+            for op in previous_ops
+        )

python/paddle/quantization/imperative/utils.py

@@ -0,0 +1,180 @@
+#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+
+import paddle
+from paddle.nn.quant import quant_layers
+
+layer_name_map = {
+    'Conv2DTranspose': paddle.nn.Conv2DTranspose,
+    'Conv2D': paddle.nn.Conv2D,
+    'Linear': paddle.nn.Linear,
+    'AdaptiveAvgPool2D': paddle.nn.AdaptiveAvgPool2D,
+    'AdaptiveMaxPool2D': paddle.nn.AdaptiveMaxPool2D,
+    'AvgPool2D': paddle.nn.AvgPool2D,
+    'MaxPool2D': paddle.nn.MaxPool2D,
+    'Hardswish': paddle.nn.Hardswish,
+    'LeakyReLU': paddle.nn.LeakyReLU,
+    'PReLU': paddle.nn.PReLU,
+    'ReLU': paddle.nn.ReLU,
+    'ReLU6': paddle.nn.ReLU6,
+    'Sigmoid': paddle.nn.Sigmoid,
+    'Softmax': paddle.nn.Softmax,
+    'Swish': paddle.nn.Swish,
+    'Tanh': paddle.nn.Tanh,
+    'BatchNorm': paddle.nn.BatchNorm,
+    'GroupNorm': paddle.nn.GroupNorm,
+    'LayerNorm': paddle.nn.LayerNorm,
+}
+
+# Apply fake quant for the inputs of these layers
+fake_quant_input_layers = [
+    paddle.nn.Conv2D,
+    paddle.nn.Linear,
+    paddle.nn.Conv2DTranspose,
+]
+
+# Apply fake quant for the output of these layers
+# TODO(jc): fix the problem of adding duplicate fake_quant ops
+# paddle.nn.AdaptiveAvgPool2D, paddle.nn.AvgPool2D, paddle.nn.ReLU,paddle.nn.LeakyReLU
+fake_quant_output_layers = [
+    paddle.nn.quant.add,
+    paddle.nn.quant.subtract,
+    paddle.nn.quant.multiply,
+    paddle.nn.quant.divide,
+    paddle.nn.quant.matmul,
+]
+
+fake_quant_leaf_layers = [
+    quant_layers.FakeQuantAbsMax,
+    quant_layers.FakeQuantChannelWiseAbsMax,
+    quant_layers.FakeQuantMovingAverageAbsMax,
+    quant_layers.MovingAverageAbsMaxScale,
+]
+
+fake_quant_wrap_layers = [
+    quant_layers.QuantizedConv2D,
+    quant_layers.QuantizedLinear,
+    quant_layers.QuantizedConv2DTranspose,
+    quant_layers.QuantizedColumnParallelLinear,
+    quant_layers.QuantizedRowParallelLinear,
+]
+
+# The weight format of these layers is Cin * Cout * H * W
+spec_channel_axis_layers = [paddle.nn.Conv2DTranspose, paddle.nn.Linear]
+
+weight_op_types = [
+    "conv2d",
+    "depthwise_conv2d",
+    "matmul",
+    "conv2d_transpose",
+    "depthwise_conv2d_transpose",
+]
+
+fake_quantize_dequantize_op_types = [
+    "fake_quantize_dequantize_abs_max",
+    "fake_channel_wise_quantize_dequantize_abs_max",
+    "fake_quantize_dequantize_moving_average_abs_max",
+]
+
+
+def load_variable_data(scope, var_name):
+    """
+    Load variable value from scope
+    """
+    var_node = scope.find_var(var_name)
+    assert var_node is not None, "Can not find " + var_name + " in the scope."
+    return np.array(var_node.get_tensor())
+
+
+def find_previous_op(block, var_name):
+    """
+    Find the previous op for the input variable.
+    """
+    for op in block.ops:
+        if var_name in op.output_arg_names:
+            return op
+    return None
+
+
+def find_next_ops(block, var_name):
+    """
+    Find all followed ops for the input variable.
+    """
+    res_ops = []
+    for op in block.ops:
+        if var_name in op.input_arg_names:
+            res_ops.append(op)
+    return res_ops
+
+
+def find_parent_layer_and_sub_name(model, name):
+    """
+    Given the model and the name of a layer, find the parent layer and
+    the sub_name of the layer.
+    For example, if name is 'block_1/convbn_1/conv_1', the parent layer is
+    'block_1/convbn_1' and the sub_name is `conv_1`.
+    Args:
+        model(paddle.nn.Layer): the model to be quantized.
+        name(string): the name of a layer
+
+    Returns:
+        parent_layer, subname
+    """
+    assert isinstance(model, paddle.nn.Layer), (
+        "The model must be the instance of paddle.nn.Layer."
+    )
+    assert len(name) > 0, "The input (name) should not be empty."
+
+    last_idx = 0
+    idx = 0
+    parent_layer = model
+    while idx < len(name):
+        if name[idx] == '.':
+            sub_name = name[last_idx:idx]
+            if hasattr(parent_layer, sub_name):
+                parent_layer = getattr(parent_layer, sub_name)
+                last_idx = idx + 1
+        idx += 1
+    sub_name = name[last_idx:idx]
+    return parent_layer, sub_name
+
+
+def program_all_ops(program):
+    """
+    Return all ops for the input program.
+    """
+    all_ops = []
+    for block in program.blocks:
+        for op in block.ops:
+            all_ops.append(op)
+    return all_ops
+
+
+def is_leaf_layer(layer):
+    """
+    Whether the layer is leaf layer.
+    """
+    return isinstance(layer, paddle.nn.Layer) and len(layer.sublayers()) == 0
+
+
+def fp_numpy_to_naive(x_np):
+    """
+    Convert numpy to float or list.
+    """
+    if x_np.size == 1:
+        return float(x_np)
+    else:
+        return x_np.tolist()

python/paddle/quantization/observers/__init__.py

@@ -0,0 +1,20 @@
+"""Observers"""
+
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .abs_max import AbsmaxObserver
+from .groupwise import GroupWiseWeightObserver
+
+__all__ = ["AbsmaxObserver", "GroupWiseWeightObserver"]

python/paddle/quantization/observers/abs_max.py

@@ -0,0 +1,95 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import paddle
+
+from ..base_observer import BaseObserver
+from ..factory import ObserverFactory
+
+
+class AbsmaxObserver(ObserverFactory):
+    r"""
+    It collects maximum absolute values of target tensor.
+
+    Args:
+        bit_length(int, optional): Number of bits to represent an quantized integer in binary.
+        dtype(str, optional): The data type of input tensor.
+        name (str, optional): This parameter is used by developers to print debugging information. \
+            For details, please refer to :ref:`api_guide_Name`. Default is None.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> from paddle.quantization import QuantConfig
+            >>> from paddle.quantization.quanters import FakeQuanterWithAbsMaxObserver
+            >>> quanter = FakeQuanterWithAbsMaxObserver(moving_rate=0.99)
+            >>> q_config = QuantConfig(activation=quanter, weight=quanter)
+    """
+
+    def __init__(self, quant_bits=8):
+        super().__init__(quant_bits=quant_bits)
+
+    def _get_class(self):
+        return AbsmaxObserverLayer
+
+
+class AbsmaxObserverLayer(BaseObserver):
+    """
+    Per-tensor abs max quantizer.
+    """
+
+    INIT_ABS_MAX = 1e-7
+
+    def __init__(self, layer, quant_bits=8):
+        super().__init__()
+        self._quant_bits = quant_bits
+        self.abs_max_val = paddle.to_tensor(AbsmaxObserverLayer.INIT_ABS_MAX)
+        self._max = None
+        self._scale = None
+        self._zero_point = None
+
+    def forward(self, input):
+        self._min, self._max = self.cal_min_max(input)
+        return input
+
+    def cal_min_max(self, inputs):
+        abs_max_val = paddle.max(paddle.abs(inputs))
+        if self._max is not None:
+            abs_max_val = paddle.maximum(
+                abs_max_val, self._max.cast(inputs.dtype)
+            )
+        return 0, abs_max_val
+
+    def bit_length(self):
+        return self._quant_bits
+
+    def quant_axis(self):
+        return -1
+
+    def cal_thresholds(self):
+        """Compute thresholds for MAX function."""
+        if self._scale is None:
+            self._scale = self._max
+        self._zero_point = paddle.zeros_like(self._scale)
+
+    def scales(self):
+        """Return output scales."""
+        if self._scale is None:
+            self.cal_thresholds()
+        return self._scale
+
+    def zero_points(self):
+        """Return output zero points."""
+        return self._zero_point

python/paddle/quantization/observers/groupwise.py

@@ -0,0 +1,114 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+
+import paddle
+
+from ..base_observer import BaseObserver
+from ..factory import ObserverFactory
+
+
+class GroupWiseWeightObserver(ObserverFactory):
+    r"""
+    It collects channel-wise maximum absolute values of target weights.
+    Args:
+        bit_length(int, optional): Number of bits to represent an quantized integer in binary.
+        dtype(str, optional): The data type of input tensor.
+        name (str, optional): This parameter is used by developers to print debugging information. \
+            For details, please refer to :ref:`api_guide_Name`. Default is None.
+    Examples:
+        .. code-block:: pycon
+
+            >>> from paddle.quantization import QuantConfig
+            >>> from paddle.quantization.quanters import AbsMaxChannelWiseWeightObserver
+            >>> quanter = AbsMaxChannelWiseWeightObserver()
+            >>> q_config = QuantConfig(activation=None, weight=quanter)
+    """
+
+    def __init__(self, quant_bits=8, group_size=128):
+        super().__init__(quant_bits=quant_bits)
+
+    def _get_class(self):
+        return GroupWiseWeightObserverLayer
+
+
+class GroupWiseWeightObserverLayer(BaseObserver):
+    def __init__(self, layer, quant_bits=8, group_size=128):
+        super().__init__()
+        self._quant_bits = quant_bits
+        self.group_size = group_size
+        self._layer = layer
+        self._max = None
+        self._scale = None
+        self._zero_point = None
+
+    def forward(self, inputs):
+        self._max = self._cal_abs_max(inputs)
+        return inputs
+
+    def _cal_abs_max(self, inputs):
+        """Use group_size to group the input, then use the
+        absmax method to calculate the scale
+        """
+        input_shape = inputs.shape
+        assert self.group_size == 64 or self.group_size == 128, (
+            "group_size only support 64 or 128"
+        )
+        assert inputs.shape[0] % self.group_size == 0, (
+            "group_size must be a factor of input channels"
+        )
+        assert len(inputs.shape) == 2, "Currently only support 2D tensor"
+        input_processed = inputs.transpose([1, 0]).reshape(
+            [input_shape[1], input_shape[0] // self.group_size, self.group_size]
+        )
+
+        abs_max_values = paddle.max(paddle.abs(input_processed), axis=2).cast(
+            "float32"
+        )
+        abs_max_values = paddle.where(
+            abs_max_values == np.float32(0), np.float32(1e-8), abs_max_values
+        )
+        abs_max_values = abs_max_values.transpose([1, 0])
+        return abs_max_values
+
+    def min_value(self) -> float:
+        return 0.0
+
+    def max_value(self) -> float:
+        return self._max
+
+    def bit_length(self):
+        return self._quant_bits
+
+    def quant_axis(self):
+        return -1
+
+    def cal_thresholds(self):
+        """Compute thresholds for MAX function."""
+        if self._scale is None:
+            self._scale = self._max
+        self._zero_point = paddle.zeros_like(self._scale)
+
+    def scales(self):
+        """Return output scales."""
+        if self._scale is None:
+            self.cal_thresholds()
+        return self._scale
+
+    def zero_points(self):
+        """Return output zero points."""
+        if self._zero_point is None:
+            self.cal_thresholds()
+        return self._zero_point

python/paddle/quantization/ptq.py

@@ -0,0 +1,130 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+import copy
+from typing import TYPE_CHECKING
+
+from paddle.distributed import fleet
+
+from .quantize import Quantization
+
+if TYPE_CHECKING:
+    from paddle.nn import Layer
+
+    from .config import QuantConfig
+
+
+class PTQ(Quantization):
+    """
+    Applying post training quantization to the model.
+    """
+
+    def __init__(self, config: QuantConfig) -> None:
+        super().__init__(config)
+
+    def _is_parallel_training(self):
+        try:
+            if fleet.worker_num() > 2:
+                return True
+            else:
+                return False
+        except Exception:  # fleet is not initialized
+            return False
+
+    def quantize(self, model: Layer, inplace: bool = False) -> Layer:
+        r"""
+        Create a model for post-training quantization.
+
+        The quantization configuration will be propagated in the model.
+        And it will insert observers into the model to collect and compute
+        quantization parameters.
+
+        Args:
+            model(Layer): The model to be quantized.
+            inplace(bool): Whether to modify the model in-place.
+
+        Return: The prepared model for post-training quantization.
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> from paddle.quantization import PTQ, QuantConfig
+                >>> from paddle.quantization.observers import AbsmaxObserver
+                >>> from paddle.vision.models import LeNet
+
+                >>> observer = AbsmaxObserver()
+                >>> q_config = QuantConfig(activation=observer, weight=observer)
+                >>> ptq = PTQ(q_config)
+                >>> model = LeNet()
+                >>> model.eval()
+                >>> quant_model = ptq.quantize(model)
+                >>> print(quant_model)
+                LeNet(
+                  (features): Sequential(
+                    (0): QuantedConv2D(
+                      (weight_quanter): AbsmaxObserverLayer()
+                      (activation_quanter): AbsmaxObserverLayer()
+                    )
+                    (1): ObserveWrapper(
+                      (_observer): AbsmaxObserverLayer()
+                      (_observed): ReLU()
+                    )
+                    (2): ObserveWrapper(
+                      (_observer): AbsmaxObserverLayer()
+                      (_observed): MaxPool2D(kernel_size=2, stride=2, padding=0, dilation=1)
+                    )
+                    (3): QuantedConv2D(
+                      (weight_quanter): AbsmaxObserverLayer()
+                      (activation_quanter): AbsmaxObserverLayer()
+                    )
+                    (4): ObserveWrapper(
+                      (_observer): AbsmaxObserverLayer()
+                      (_observed): ReLU()
+                    )
+                    (5): ObserveWrapper(
+                      (_observer): AbsmaxObserverLayer()
+                      (_observed): MaxPool2D(kernel_size=2, stride=2, padding=0, dilation=1)
+                    )
+                  )
+                  (fc): Sequential(
+                    (0): QuantedLinear(
+                      (weight_quanter): AbsmaxObserverLayer()
+                      (activation_quanter): AbsmaxObserverLayer()
+                    )
+                    (1): QuantedLinear(
+                      (weight_quanter): AbsmaxObserverLayer()
+                      (activation_quanter): AbsmaxObserverLayer()
+                    )
+                    (2): QuantedLinear(
+                      (weight_quanter): AbsmaxObserverLayer()
+                      (activation_quanter): AbsmaxObserverLayer()
+                    )
+                  )
+                )
+        """
+        _model = model
+        if not inplace:
+            assert not self._is_parallel_training(), (
+                "'inplace' is not compatible with parallel training."
+            )
+            _model = copy.deepcopy(model)
+            _model.eval()
+        assert not model.training, (
+            "Post-Training Quantization should not work on training models. Please set evaluation mode by model.eval()."
+        )
+        self._config._specify(_model)
+        self._convert_to_quant_layers(_model, self._config)
+        self._insert_activation_observers(_model, self._config)
+        return _model

python/paddle/quantization/qat.py

@@ -0,0 +1,122 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+import copy
+from typing import TYPE_CHECKING
+
+from .quantize import Quantization
+
+if TYPE_CHECKING:
+    from paddle.nn import Layer
+
+    from .config import QuantConfig
+
+
+class QAT(Quantization):
+    r"""
+    Tools used to prepare model for quantization-aware training.
+    Args:
+        config(QuantConfig): Quantization configuration
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> from paddle.quantization import QAT, QuantConfig
+            >>> from paddle.quantization.quanters import FakeQuanterWithAbsMaxObserver
+            >>> quanter = FakeQuanterWithAbsMaxObserver(moving_rate=0.9)
+            >>> q_config = QuantConfig(activation=quanter, weight=quanter)
+            >>> qat = QAT(q_config)
+    """
+
+    def __init__(self, config: QuantConfig) -> None:
+        super().__init__(config)
+
+    def quantize(self, model: Layer, inplace: bool = False) -> Layer:
+        r"""
+        Create a model for quantization-aware training.
+
+        The quantization configuration will be propagated in the model.
+        And it will insert fake quanters into the model to simulate the quantization.
+
+        Args:
+            model(Layer): The model to be quantized.
+            inplace(bool): Whether to modify the model in-place.
+
+        Return: The prepared model for quantization-aware training.
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> from paddle.quantization import QAT, QuantConfig
+                >>> from paddle.quantization.quanters import FakeQuanterWithAbsMaxObserver
+                >>> from paddle.vision.models import LeNet
+
+                >>> quanter = FakeQuanterWithAbsMaxObserver(moving_rate=0.9)
+                >>> q_config = QuantConfig(activation=quanter, weight=quanter)
+                >>> qat = QAT(q_config)
+                >>> model = LeNet()
+                >>> quant_model = qat.quantize(model)
+                >>> print(quant_model)
+                LeNet(
+                  (features): Sequential(
+                    (0): QuantedConv2D(
+                      (weight_quanter): FakeQuanterWithAbsMaxObserverLayer()
+                      (activation_quanter): FakeQuanterWithAbsMaxObserverLayer()
+                    )
+                    (1): ObserveWrapper(
+                      (_observer): FakeQuanterWithAbsMaxObserverLayer()
+                      (_observed): ReLU()
+                    )
+                    (2): ObserveWrapper(
+                      (_observer): FakeQuanterWithAbsMaxObserverLayer()
+                      (_observed): MaxPool2D(kernel_size=2, stride=2, padding=0, dilation=1)
+                    )
+                    (3): QuantedConv2D(
+                      (weight_quanter): FakeQuanterWithAbsMaxObserverLayer()
+                      (activation_quanter): FakeQuanterWithAbsMaxObserverLayer()
+                    )
+                    (4): ObserveWrapper(
+                      (_observer): FakeQuanterWithAbsMaxObserverLayer()
+                      (_observed): ReLU()
+                    )
+                    (5): ObserveWrapper(
+                      (_observer): FakeQuanterWithAbsMaxObserverLayer()
+                      (_observed): MaxPool2D(kernel_size=2, stride=2, padding=0, dilation=1)
+                    )
+                  )
+                  (fc): Sequential(
+                    (0): QuantedLinear(
+                      (weight_quanter): FakeQuanterWithAbsMaxObserverLayer()
+                      (activation_quanter): FakeQuanterWithAbsMaxObserverLayer()
+                    )
+                    (1): QuantedLinear(
+                      (weight_quanter): FakeQuanterWithAbsMaxObserverLayer()
+                      (activation_quanter): FakeQuanterWithAbsMaxObserverLayer()
+                    )
+                    (2): QuantedLinear(
+                      (weight_quanter): FakeQuanterWithAbsMaxObserverLayer()
+                      (activation_quanter): FakeQuanterWithAbsMaxObserverLayer()
+                    )
+                  )
+                )
+        """
+        assert model.training, (
+            "Quantization-Aware Training should work on training models. Please set training mode by model.train()."
+        )
+        _model = model if inplace else copy.deepcopy(model)
+        self._config._specify(_model)
+        self._convert_to_quant_layers(_model, self._config)
+        self._insert_activation_observers(_model, self._config)
+        return _model

python/paddle/quantization/quanters/__init__.py

@@ -0,0 +1,17 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .abs_max import FakeQuanterWithAbsMaxObserver
+
+__all__ = ["FakeQuanterWithAbsMaxObserver"]

python/paddle/quantization/quanters/abs_max.py

@@ -0,0 +1,269 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import paddle
+from paddle import _C_ops
+from paddle.base.data_feeder import check_variable_and_dtype
+from paddle.base.framework import _create_tensor
+from paddle.framework import ParamAttr, core
+from paddle.nn.initializer import Constant
+from paddle.utils import unique_name
+
+from ..base_quanter import BaseQuanter
+from ..factory import QuanterFactory
+
+
+class FakeQuanterWithAbsMaxObserver(QuanterFactory):
+    r"""
+    Compute quantization parameters and simulate quantization.
+
+    It collects maximum absolute values of target tensor with moving average.
+    The average value will be used as quantization scale to quantize and
+    dequantize the tensor.
+
+    And it is symmetric uniform quantization which means the zero point is always 0.
+
+    The computational formula of moving average is described as below:
+
+    .. math::
+            state = rate * state + 1
+            accum = rate * accum + max(abs(x))
+            scale = accum / state
+
+    Where:
+
+    - :math:`x` is the input tensor.
+    - :math:`state` and :math:`accum` are zero-initialized accumulators.
+    - :math:`rate` is moving average rate.
+    - :math:`scale` is quantization scale
+
+    And the computational formula of simulate quantization is:
+
+    .. math::
+            range = 2^{bit\_length - 1} - 1
+            out = round(x / scale * range) * scale / range
+
+    Where:
+
+    - :math:`{bit\_length}` is the length of bits.
+    - :math:`x` is the input tensor and :math:`out` is the output of simulate quantization.
+
+    Args:
+        moving_rate(float, optional): The rate of moving average.
+        bit_length(int, optional): Number of bits to represent an quantized integer in binary.
+        dtype(str, optional): The data type of input tensor.
+        name (str, optional): This parameter is used by developers to print debugging information. \
+            For details, please refer to :ref:`api_guide_Name`. Default is None.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> from paddle.quantization import QuantConfig
+            >>> from paddle.quantization.quanters import FakeQuanterWithAbsMaxObserver
+            >>> quanter = FakeQuanterWithAbsMaxObserver(moving_rate=0.99)
+            >>> q_config = QuantConfig(activation=quanter, weight=quanter)
+    """
+
+    def __init__(
+        self,
+        moving_rate=0.9,
+        bit_length=8,
+        dtype='float32',
+        name=None,
+    ):
+        super().__init__(
+            name=name,
+            moving_rate=moving_rate,
+            bit_length=bit_length,
+            dtype=dtype,
+        )
+
+    def _get_class(self):
+        return FakeQuanterWithAbsMaxObserverLayer
+
+
+class FakeQuanterWithAbsMaxObserverLayer(BaseQuanter):
+    def __init__(
+        self,
+        layer,
+        name=None,
+        moving_rate=0.9,
+        bit_length=8,
+        dtype='float32',
+    ):
+        super().__init__()
+        self._moving_rate = moving_rate
+        self._bit_length = bit_length
+        scale_prefix = f"{name}.scale" if name else 'quant_dequant.scale'
+        scale_attr = ParamAttr(
+            name=unique_name.generate(scale_prefix),
+            initializer=Constant(0.001),
+            trainable=False,
+        )
+        self._scale = self.create_parameter(
+            shape=[1], attr=scale_attr, dtype=dtype
+        )
+        self._scale.stop_gradient = True
+
+        state_prefix = f"{name}.state" if name else 'quant_dequant.state'
+        state_attr = ParamAttr(
+            name=unique_name.generate(state_prefix),
+            initializer=Constant(1),
+            trainable=False,
+        )
+        self._state = self.create_parameter(
+            shape=[1], attr=state_attr, dtype=dtype
+        )
+        self._state.stop_gradient = True
+
+        accum_prefix = f"{name}.accum" if name else 'quant_dequant.accum'
+        accum_attr = ParamAttr(
+            name=unique_name.generate(accum_prefix),
+            initializer=Constant(1),
+            trainable=False,
+        )
+        self._accum = self.create_parameter(
+            shape=[1], attr=accum_attr, dtype=dtype
+        )
+        self._accum.stop_gradient = True
+
+    def dynamic_forward(self, input):
+        attrs = (
+            'moving_rate',
+            self._moving_rate,
+            'bit_length',
+            self._bit_length,
+            'is_test',
+            not self.training,
+        )
+        quant_out = _create_tensor(
+            type=input.type,
+            name=f"{input.name}.quantized.dequantized",
+            shape=input.shape,
+            dtype=input.dtype,
+            persistable=False,
+        )
+
+        state = self._state if self.training else None
+        accum = self._accum if self.training else None
+
+        (
+            out1,
+            out2,
+            out3,
+            out4,
+        ) = _C_ops.fake_quantize_dequantize_moving_average_abs_max(
+            input,
+            self._scale,
+            accum,
+            state,
+            self._moving_rate,
+            self._bit_length,
+            not self.training,
+            1,
+        )
+        _C_ops.assign_out_(out1, quant_out)
+        if out2._is_initialized():
+            _C_ops.assign_out_(out2, self._scale)
+        if state:
+            _C_ops.assign_out_(out3, state)
+        if accum:
+            _C_ops.assign_out_(out4, accum)
+        return quant_out
+
+    def static_forward(self, input):
+        check_variable_and_dtype(
+            input, 'input', ['float32'], "FakeQuantMovingAverageAbsMax"
+        )
+        attrs = {
+            'moving_rate': self._moving_rate,
+            'bit_length': self._bit_length,
+            'is_test': not self.training,
+        }
+        inputs = {"X": [input], "InScale": [self._scale]}
+        quant_out = self._helper.create_variable(
+            name=f"{input.name}.quantized.dequantized",
+            dtype=input.dtype,
+            type=core.VarDesc.VarType.DENSE_TENSOR,
+            persistable=False,
+            stop_gradient=False,
+        )
+        outputs = {"Out": [quant_out], "OutScale": [self._scale]}
+
+        if self.training:
+            inputs['InState'] = [self._state]
+            inputs['InAccum'] = [self._accum]
+            outputs['OutState'] = [self._state]
+            outputs['OutAccum'] = [self._accum]
+
+        self._helper.append_op(
+            type="fake_quantize_dequantize_moving_average_abs_max",
+            inputs=inputs,
+            outputs=outputs,
+            attrs=attrs,
+        )
+
+        return quant_out
+
+    def pir_forward(self, input):
+        state = self._state if self.training else None
+        accum = self._accum if self.training else None
+
+        (
+            out1,
+            out2,
+            out3,
+            out4,
+        ) = _C_ops.fake_quantize_dequantize_moving_average_abs_max(
+            input,
+            self._scale,
+            accum,
+            state,
+            self._moving_rate,
+            self._bit_length,
+            not self.training,
+            1,
+        )
+
+        # TODO, need to check this modify can work correctly in PIR mode
+        # there is no `name` attribute of Value in PIR mode
+        # so, directly return quant_out in pir mode
+        quant_out = out1
+        _C_ops.assign_out_(out2, self._scale)
+
+        if self.training:
+            _C_ops.assign_out_(out3, state)
+        if self.training:
+            _C_ops.assign_out_(out4, accum)
+        return quant_out
+
+    def forward(self, input):
+        if paddle.in_dynamic_mode():
+            return self.dynamic_forward(input)
+        elif paddle.base.framework.in_pir_mode():
+            return self.pir_forward(input)
+        else:
+            return self.static_forward(input)
+
+    def bit_length(self):
+        return self._bit_length
+
+    def quant_axis(self):
+        return -1
+
+    def scales(self):
+        return self._scale
+
+    def zero_points(self):
+        return None

python/paddle/quantization/quantize.py

@@ -0,0 +1,128 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import abc
+import copy
+
+from paddle.nn import Layer
+from paddle.nn.quant.format import (
+    ConvertibleQuantedLayer,
+    LinearQuanterDequanter,
+)
+
+from .base_quanter import BaseQuanter
+from .config import QuantConfig
+
+
+class Quantization(metaclass=abc.ABCMeta):
+    r"""
+    Abstract class used to prepares a copy of the model for quantization calibration or quantization-aware training.
+    Args:
+        config(QuantConfig): Quantization configuration
+    """
+
+    def __init__(self, config: QuantConfig):
+        self._config = copy.deepcopy(config)
+
+    @abc.abstractmethod
+    def quantize(self, model: Layer, inplace=False):
+        r"""Create a model for quantization-aware training or post-training quantization."""
+        pass
+
+    def convert(self, model: Layer, inplace=False, remain_weight=False):
+        r"""Convert the quantization model to ONNX style. And the converted
+        model can be saved as inference model by calling paddle.jit.save.
+        Args:
+            model(Layer): The quantized model to be converted.
+            inplace(bool, optional): Whether to modify the model in-place, default is False.
+            remain_weight(bool, optional): Whether to remain weights in floats, default is False.
+
+        Return: The converted model
+
+        Examples:
+            .. code-block:: pycon
+
+                >>> import paddle
+                >>> from paddle.quantization import QAT, QuantConfig
+                >>> from paddle.quantization.quanters import FakeQuanterWithAbsMaxObserver
+                >>> from paddle.vision.models import LeNet
+
+                >>> quanter = FakeQuanterWithAbsMaxObserver(moving_rate=0.9)
+                >>> q_config = QuantConfig(activation=quanter, weight=quanter)
+                >>> qat = QAT(q_config)
+                >>> model = LeNet()
+                >>> quantized_model = qat.quantize(model)
+                >>> converted_model = qat.convert(quantized_model)
+                >>> dummy_data = paddle.rand([1, 1, 32, 32], dtype="float32")
+                >>> paddle.jit.save(converted_model, "./quant_deploy", [dummy_data])
+        """
+        _model = model if inplace else copy.deepcopy(model)
+        replaced = {}
+        for name, child in _model.named_children():
+            quant_dequant = None
+            if isinstance(child, ConvertibleQuantedLayer):
+                if child.converted:
+                    continue
+                if hasattr(child, 'weight_quanter') and (
+                    child.weight_quanter is None
+                    or child.weight_quanter.scales() is None
+                ):
+                    continue
+                child._convert(remain_weight=remain_weight)
+            elif isinstance(child, BaseQuanter):
+                quant_dequant = LinearQuanterDequanter.from_quanter(child)
+            else:
+                self.convert(child, inplace=True, remain_weight=remain_weight)
+            if quant_dequant is not None:
+                replaced[name] = quant_dequant
+        for key, value in replaced.items():
+            _model._sub_layers[key] = value
+        return _model
+
+    def _convert_to_quant_layers(self, model: Layer, config: QuantConfig):
+        replaced = {}
+        for name, child in model.named_children():
+            if (
+                config._is_quantifiable(child)
+                and type(child) in config.qat_layer_mappings
+            ):
+                replaced[name] = config._get_qat_layer(child)
+            else:
+                self._convert_to_quant_layers(child, config)
+        for key, value in replaced.items():
+            model._sub_layers[key] = value
+
+    def _insert_activation_observers(self, model: Layer, config: QuantConfig):
+        replaced = {}
+        for name, child in model.named_children():
+            if config._need_observe(child):
+                replaced[name] = config._get_observe_wrapper(child)
+            else:
+                if (
+                    type(child) not in config._qat_layer_mapping.values()
+                    and type(child)
+                    not in config._customized_qat_layer_mapping.values()
+                ):
+                    self._insert_activation_observers(child, config)
+        for key, value in replaced.items():
+            model._sub_layers[key] = value
+
+    def _details(self):
+        return self._config.details()
+
+    def __str__(self):
+        return self._details()
+
+    def __repr__(self):
+        return self.__str__()

python/paddle/quantization/wrapper.py

@@ -0,0 +1,48 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from paddle.nn import Layer
+
+from .base_quanter import BaseQuanter
+
+
+class ObserveWrapper(Layer):
+    r"""
+    Put an observer layer and an observed layer into a wrapping layer.
+    It is used to insert layers into the model for QAT or PTQ.
+    Args:
+        observer(BaseQuanter): Observer layer
+        observed(Layer): Observed layer
+        observe_input(bool): If it is true the observer layer will be called before observed layer.
+            If it is false the observed layer will be called before observer layer. Default: True.
+    """
+
+    def __init__(
+        self,
+        observer: BaseQuanter,
+        observed: Layer,
+        observe_input=True,
+    ):
+        super().__init__()
+        self._observer = observer
+        self._observed = observed
+        self._observe_input = observe_input
+
+    def forward(self, *inputs, **kwargs):
+        if self._observe_input:
+            out = self._observer(*inputs, **kwargs)
+            return self._observed(out, **kwargs)
+        else:
+            out = self._observed(*inputs, **kwargs)
+            return self._observer(out, **kwargs)