chore: import upstream snapshot with attribution

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]