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