Paolify/RVC_4

1# Copyright (c) Facebook, Inc. and its affiliates.2# All rights reserved.3#4# This source code is licensed under the license found in the5# LICENSE file in the root directory of this source tree.6 7"""8Classic uniform quantization over n bits.9"""10from typing import Tuple11import torch12 13from .base import BaseQuantizer14from .utils import simple_repr15 16 17def uniform_quantize(p: torch.Tensor, bits: torch.Tensor = torch.tensor(8.)):18    """19    Quantize the given weights over `bits` bits.20 21    Returns:22        - quantized levels23        - (min, max) range.24 25    """26    assert (bits >= 1).all() and (bits <= 15).all()27    num_levels = (2 ** bits.float()).long()28    mn = p.min().item()29    mx = p.max().item()30    p = (p - mn) / (mx - mn)  # put p in [0, 1]31    unit = 1 / (num_levels - 1)  # quantization unit32    levels = (p / unit).round()33    if (bits <= 8).all():34        levels = levels.byte()35    else:36        levels = levels.short()37    return levels, (mn, mx)38 39 40def uniform_unquantize(levels: torch.Tensor, scales: Tuple[float, float],41                       bits: torch.Tensor = torch.tensor(8.)):42    """43    Unquantize the weights from the levels and scale. Return a float32 tensor.44    """45    mn, mx = scales46    num_levels = 2 ** bits.float()47    unit = 1 / (num_levels - 1)48    levels = levels.float()49    p = levels * unit  # in [0, 1]50    return p * (mx - mn) + mn51 52 53class UniformQuantizer(BaseQuantizer):54    def __init__(self, model: torch.nn.Module, bits: float = 8., min_size: float = 0.01,55                 float16: bool = False, qat: bool = False, exclude=[], detect_bound=True):56        """57        Args:58            model (torch.nn.Module): model to quantize59            bits (float): number of bits to quantize over.60            min_size (float): minimum size in MB of a parameter to be quantized.61            float16 (bool): if a layer is smaller than min_size, should we still do float16?62            qat (bool): perform quantized aware training.63            exclude (list[str]): list of patterns used to match parameters to exclude.64                For instance `['bias']` to exclude all bias terms.65            detect_bound (bool): if True, will detect bound parameters and reuse66                the same quantized tensor for both.67        """68        self.bits = float(bits)69        self.qat = qat70 71        super().__init__(model, min_size, float16, exclude, detect_bound)72 73    def __repr__(self):74        return simple_repr(self, )75 76    def _pre_forward_train(self):77        if self.qat:78            for qparam in self._qparams:79                if qparam.other is not None:80                    new_param = qparam.other.module._parameters[qparam.other.name]81                else:82                    quantized = self._quantize_param(qparam)83                    qvalue = self._unquantize_param(qparam, quantized)84                    new_param = qparam.param + (qvalue - qparam.param).detach()85                qparam.module._parameters[qparam.name] = new_param86            return True87        return False88 89    def _post_forward_train(self):90        if self.qat:91            for qparam in self._qparams:92                qparam.module._parameters[qparam.name] = qparam.param93            return True94        return False95 96    def _quantize_param(self, qparam):97        levels, scales = uniform_quantize(qparam.param.data, torch.tensor(self.bits))98        return (levels, scales)99 100    def _unquantize_param(self, qparam, quantized):101        levels, scales = quantized102        return uniform_unquantize(levels, scales, torch.tensor(self.bits))103 104    def model_size(self):105        """106        Non differentiable model size in MB.107        """108        total = super().model_size()109        subtotal = 0110        for qparam in self._qparams:111            if qparam.other is None:  # if parameter is bound, count only one copy.112                subtotal += self.bits * qparam.param.numel() + 64  # 2 float for the overall scales113        subtotal /= 2**20 * 8  # bits to MegaBytes114        return total + subtotal115 116    def true_model_size(self):117        """118        Return the true quantized model size, in MB, without extra119        compression.120        """121        return self.model_size().item()122