Spaces:
V3N0M
/
Runtime error

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()