Source code for tensorplay.ao.quantization.stubs

"""Quant/DeQuant stubs for float<->quantized conversion points in a model.

QuantStub calibrates on observed activations (training/QAT) and, once
frozen, converts incoming floats to native QInt8 tensors via the
quantize_per_tensor kernel.  DeQuantStub converts a quantized tensor back to
Float32 through its own quantizer (or explicit parameters for raw code
tensors).
"""

import tensorplay
from tensorplay._C import (
    _make_per_tensor_quantized_tensor as _make_per_tensor_quantized_tensor,
    quantize_per_tensor as _quantize_per_tensor,
)
from tensorplay import nn

from .fake_quantize import FakeQuantize

__all__ = ["QuantStub", "DeQuantStub", "QuantWrapper"]



[docs]
class QuantStub(nn.Module):
    def __init__(self, qconfig=None):
        super().__init__()
        if qconfig is not None and hasattr(qconfig, 'activation'):
            self.qconfig = qconfig
            self.fake_quant = FakeQuantize(observer=qconfig.activation())
        else:
            self.fake_quant = FakeQuantize() if qconfig is None else qconfig()


[docs]
    def record(self, x):
        """Calibration entry point: feeds the batch to the inner FakeQuantize
        observer without fake-quantizing (as manual calibration loops do)."""
        self.fake_quant.record(x)
        return x


    def forward(self, x):
        # Calibration / QAT path: simulated quantization keeps the graph
        # float while nudging values toward the quantized grid.
        if hasattr(self, 'activation_post_process'):
            return x
        scale, zero_point = self.fake_quant.scale, self.fake_quant.zero_point
        if not self.training and self.fake_quant.frozen:
            # Inference path: produce a native quantized tensor carrying
            # its affine parameters.
            return _quantize_per_tensor(self=x, scale=float(scale),
                                        zero_point=int(zero_point),
                                        dtype=self.fake_quant.dtype)
        return self.fake_quant(x)

    def freeze(self):
        self.fake_quant.freeze()




[docs]
class DeQuantStub(nn.Module):
    def __init__(self, scale=None, zero_point=None):
        super().__init__()
        self.scale = scale
        self.zero_point = zero_point

    def forward(self, x):
        if x.is_quantized():
            # Native path: the tensor carries its own affine parameters.
            return x.dequantize()
        if x.dtype != tensorplay.int8:
            # Before conversion the surrounding model still runs in float, so
            # the stub stays an identity until it is swapped for a real
            # dequantizer.
            return x
        scale = 1.0 if self.scale is None else float(self.scale)
        zero_point = 0 if self.zero_point is None else int(self.zero_point)
        q = _make_per_tensor_quantized_tensor(x, scale, zero_point)
        return q.dequantize()



class QuantWrapper(nn.Module):
    """Wrap a module with quant/dequant stubs on either side.

    Before conversion the stubs behave as observers/identity, so the wrapped
    module can be calibrated in place; after conversion the stubs become the
    native quantize/dequantize operators.
    """

    def __init__(self, module):
        super().__init__()
        qconfig = getattr(module, "qconfig", None)
        self.add_module("quant", QuantStub(qconfig))
        self.add_module("dequant", DeQuantStub())
        self.add_module("module", module)
        self.train(module.training)

    def forward(self, x):
        x = self.quant(x)
        x = self.module(x)
        return self.dequant(x)