|
|
|
|
|
|
|
|
from torch import nn |
|
|
|
|
|
|
|
|
class QuantStub(nn.Module): |
|
|
r"""Quantize stub module, before calibration, this is same as an observer, |
|
|
it will be swapped as `nnq.Quantize` in `convert`. |
|
|
|
|
|
Args: |
|
|
qconfig: quantization configuration for the tensor, |
|
|
if qconfig is not provided, we will get qconfig from parent modules |
|
|
""" |
|
|
|
|
|
def __init__(self, qconfig=None): |
|
|
super().__init__() |
|
|
if qconfig: |
|
|
self.qconfig = qconfig |
|
|
|
|
|
def forward(self, x): |
|
|
return x |
|
|
|
|
|
|
|
|
class DeQuantStub(nn.Module): |
|
|
r"""Dequantize stub module, before calibration, this is same as identity, |
|
|
this will be swapped as `nnq.DeQuantize` in `convert`. |
|
|
|
|
|
Args: |
|
|
qconfig: quantization configuration for the tensor, |
|
|
if qconfig is not provided, we will get qconfig from parent modules |
|
|
""" |
|
|
|
|
|
def __init__(self, qconfig=None): |
|
|
super().__init__() |
|
|
if qconfig: |
|
|
self.qconfig = qconfig |
|
|
|
|
|
def forward(self, x): |
|
|
return x |
|
|
|
|
|
|
|
|
class QuantWrapper(nn.Module): |
|
|
r"""A wrapper class that wraps the input module, adds QuantStub and |
|
|
DeQuantStub and surround the call to module with call to quant and dequant |
|
|
modules. |
|
|
|
|
|
This is used by the `quantization` utility functions to add the quant and |
|
|
dequant modules, before `convert` function `QuantStub` will just be observer, |
|
|
it observes the input tensor, after `convert`, `QuantStub` |
|
|
will be swapped to `nnq.Quantize` which does actual quantization. Similarly |
|
|
for `DeQuantStub`. |
|
|
""" |
|
|
quant: QuantStub |
|
|
dequant: DeQuantStub |
|
|
module: nn.Module |
|
|
|
|
|
def __init__(self, module): |
|
|
super().__init__() |
|
|
qconfig = getattr(module, "qconfig", None) |
|
|
self.add_module("quant", QuantStub(qconfig)) |
|
|
self.add_module("dequant", DeQuantStub(qconfig)) |
|
|
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) |
|
|
|
