mosi77/zan

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 7from dataclasses import dataclass8from concurrent import futures9from fnmatch import fnmatch10from functools import partial11import io12import math13from multiprocessing import cpu_count14import typing as tp15import zlib16 17import torch18 19 20class BaseQuantizer:21    @dataclass22    class _QuantizedParam:23        name: str24        param: torch.nn.Parameter25        module: torch.nn.Module26        # If a Parameter is used multiple times, `other` can be used27        # to share state between the different Quantizers28        other: tp.Optional[tp.Any]29 30    def __init__(self, model: torch.nn.Module, min_size: float = 0.01, float16: bool = False,31                 exclude: tp.Optional[tp.List[str]] = [], detect_bound: bool = True):32        self.model = model33        self.min_size = min_size34        self.float16 = float1635        self.exclude = exclude36        self.detect_bound = detect_bound37        self._quantized = False38        self._pre_handle = self.model.register_forward_pre_hook(self._forward_pre_hook)39        self._post_handle = self.model.register_forward_hook(self._forward_hook)40 41        self._quantized_state = None42        self._qparams = []43        self._float16 = []44        self._others = []45        self._rnns = []46 47        self._saved = []48 49        self._find_params()50 51    def _find_params(self):52        min_params = self.min_size * 2**20 // 453        previous = {}54        for module_name, module in self.model.named_modules():55            if isinstance(module, torch.nn.RNNBase):56                self._rnns.append(module)57            for name, param in list(module.named_parameters(recurse=False)):58                full_name = f"{module_name}.{name}"59                matched = False60                for pattern in self.exclude:61                    if fnmatch(full_name, pattern) or fnmatch(name, pattern):62                        matched = True63                        break64 65                if param.numel() <= min_params or matched:66                    if id(param) in previous:67                        continue68                    if self.detect_bound:69                        previous[id(param)] = None70                    if self.float16:71                        self._float16.append(param)72                    else:73                        self._others.append(param)74                else:75                    qparam = self._register_param(name, param, module, previous.get(id(param)))76                    if self.detect_bound:77                        previous[id(param)] = qparam78                    self._qparams.append(qparam)79 80    def _register_param(self, name, param, module, other):81        return self.__class__._QuantizedParam(name, param, module, other)82 83    def _forward_pre_hook(self, module, input):84        if self.model.training:85            self._quantized_state = None86            if self._quantized:87                self.unquantize()88            if self._pre_forward_train():89                self._fix_rnns()90        else:91            self.quantize()92 93    def _forward_hook(self, module, input, output):94        if self.model.training:95            if self._post_forward_train():96                self._fix_rnns(flatten=False)  # Hacky, next forward will flatten97 98    def quantize(self, save=True):99        """100        Immediately apply quantization to the model parameters.101        If `save` is True, save a copy of the unquantized parameters, that can be102        restored with `unquantize()`.103        """104        if self._quantized:105            return106        if save:107            self._saved = [qp.param.data.to('cpu', copy=True)108                           for qp in self._qparams if qp.other is None]109        self.restore_quantized_state(self.get_quantized_state())110        self._quantized = True111        self._fix_rnns()112 113    def unquantize(self):114        """115        Revert a previous call to `quantize()`.116        """117        if not self._quantized:118            raise RuntimeError("Can only be called on a quantized model.")119        if not self._saved:120            raise RuntimeError("Nothing to restore.")121        for qparam in self._qparams:122            if qparam.other is None:123                qparam.param.data[:] = self._saved.pop(0)124        assert len(self._saved) == 0125        self._quantized = False126        self._fix_rnns()127 128    def _pre_forward_train(self) -> bool:129        """130        Called once before each forward for continuous quantization.131        Should return  True if parameters were changed.132        """133        return False134 135    def _post_forward_train(self) -> bool:136        """137        Called once after each forward (to restore state for instance).138        Should return True if parameters were changed.139        """140        return False141 142    def _fix_rnns(self, flatten=True):143        """144        To be called after quantization happened to fix RNNs.145        """146        for rnn in self._rnns:147            rnn._flat_weights = [148                (lambda wn: getattr(rnn, wn) if hasattr(rnn, wn) else None)(wn)149                for wn in rnn._flat_weights_names]150            if flatten:151                rnn.flatten_parameters()152 153    def get_quantized_state(self):154        """155        Returns sufficient quantized information to rebuild the model state.156 157        ..Note::158            To achieve maximum compression, you should compress this with159            gzip or other, as quantized weights are not optimally coded!160        """161        if self._quantized_state is None:162            self._quantized_state = self._get_quantized_state()163        return self._quantized_state164 165    def _get_quantized_state(self):166        """167        Actual implementation for `get_quantized_state`.168        """169        float16_params = []170        for p in self._float16:171            q = p.data.half()172            float16_params.append(q)173 174        return {175            "quantized": [self._quantize_param(qparam) for qparam in self._qparams176                          if qparam.other is None],177            "float16": float16_params,178            "others": [p.data.clone() for p in self._others],179        }180 181    def _quantize_param(self, qparam: _QuantizedParam) -> tp.Any:182        """183        To be overriden.184        """185        raise NotImplementedError()186 187    def _unquantize_param(self, qparam: _QuantizedParam, quantized: tp.Any) -> torch.Tensor:188        """189        To be overriden.190        """191        raise NotImplementedError()192 193    def restore_quantized_state(self, state) -> None:194        """195        Restore the state of the model from the quantized state.196        """197        for p, q in zip(self._float16, state["float16"]):198            p.data[:] = q.to(p)199 200        for p, q in zip(self._others, state["others"]):201            p.data[:] = q202 203        remaining = list(state["quantized"])204        for qparam in self._qparams:205            if qparam.other is not None:206                # Only unquantize first appearance of nn.Parameter.207                continue208            quantized = remaining.pop(0)209            qparam.param.data[:] = self._unquantize_param(qparam, quantized)210        self._fix_rnns()211 212    def detach(self) -> None:213        """214        Detach from the model, removes hooks and anything else.215        """216        self._pre_handle.remove()217        self._post_handle.remove()218 219    def model_size(self) -> torch.Tensor:220        """221        Returns an estimate of the quantized model size.222        """223        total = torch.tensor(0.)224        for p in self._float16:225            total += 16 * p.numel()226        for p in self._others:227            total += 32 * p.numel()228        return total / 2**20 / 8  # bits to MegaBytes229 230    def true_model_size(self) -> float:231        """232        Return the true quantized model size, in MB, without extra233        compression.234        """235        return self.model_size().item()236 237    def compressed_model_size(self, compress_level=-1, num_workers=8) -> float:238        """239        Return the compressed quantized model size, in MB.240 241        Args:242            compress_level (int): compression level used with zlib,243                see `zlib.compress` for details.244            num_workers (int): will split the final big byte representation in that245                many chunks processed in parallels.246        """247        out = io.BytesIO()248        torch.save(self.get_quantized_state(), out)249        ms = _parallel_compress_len(out.getvalue(), compress_level, num_workers)250        return ms / 2 ** 20251 252 253def _compress_len(data, compress_level):254    return len(zlib.compress(data, level=compress_level))255 256 257def _parallel_compress_len(data, compress_level, num_workers):258    num_workers = min(cpu_count(), num_workers)259    chunk_size = int(math.ceil(len(data) / num_workers))260    chunks = [data[offset:offset + chunk_size] for offset in range(0, len(data), chunk_size)]261    with futures.ProcessPoolExecutor(num_workers) as pool:262        return sum(pool.map(partial(_compress_len, compress_level=compress_level), chunks))263