270 lines
8.2 KiB
Python
270 lines
8.2 KiB
Python
# 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
|