|
|
|
|
|
|
|
|
import torch |
|
|
from torch.ao.quantization.qconfig import QConfig |
|
|
from torch.ao.quantization.quant_type import QuantType |
|
|
from torch.jit._recursive import wrap_cpp_module |
|
|
|
|
|
|
|
|
__all__ = [ |
|
|
"script_qconfig", |
|
|
"script_qconfig_dict", |
|
|
"fuse_conv_bn_jit", |
|
|
"prepare_jit", |
|
|
"prepare_dynamic_jit", |
|
|
"convert_jit", |
|
|
"convert_dynamic_jit", |
|
|
"quantize_jit", |
|
|
"quantize_dynamic_jit", |
|
|
] |
|
|
|
|
|
|
|
|
def _check_is_script_module(model): |
|
|
if not isinstance(model, torch.jit.ScriptModule): |
|
|
raise ValueError("input must be a script module, got: " + str(type(model))) |
|
|
|
|
|
|
|
|
def _check_forward_method(model): |
|
|
if not model._c._has_method("forward"): |
|
|
raise ValueError("input script module does not have forward method") |
|
|
|
|
|
|
|
|
def script_qconfig(qconfig): |
|
|
r"""Instantiate the activation and weight observer modules and script |
|
|
them, these observer module instances will be deepcopied during |
|
|
prepare_jit step. |
|
|
""" |
|
|
return QConfig( |
|
|
activation=torch.jit.script(qconfig.activation())._c, |
|
|
weight=torch.jit.script(qconfig.weight())._c, |
|
|
) |
|
|
|
|
|
|
|
|
def script_qconfig_dict(qconfig_dict): |
|
|
r"""Helper function used by `prepare_jit`. |
|
|
Apply `script_qconfig` for all entries in `qconfig_dict` that is |
|
|
not None. |
|
|
""" |
|
|
return {k: script_qconfig(v) if v else None for k, v in qconfig_dict.items()} |
|
|
|
|
|
|
|
|
def fuse_conv_bn_jit(model, inplace=False): |
|
|
r"""Fuse conv - bn module |
|
|
Works for eval model only. |
|
|
|
|
|
Args: |
|
|
model: TorchScript model from scripting or tracing |
|
|
""" |
|
|
torch._C._log_api_usage_once("quantization_api.quantize_jit.fuse_conv_bn_jit") |
|
|
model_c = model._c |
|
|
model_c = torch._C._jit_pass_fold_convbn(model_c) |
|
|
if inplace: |
|
|
model._reconstruct(model_c) |
|
|
else: |
|
|
model = wrap_cpp_module(model_c) |
|
|
return model |
|
|
|
|
|
|
|
|
def _prepare_jit(model, qconfig_dict, inplace=False, quant_type=QuantType.STATIC): |
|
|
_check_is_script_module(model) |
|
|
_check_forward_method(model) |
|
|
if not all(isinstance(x, str) for x in qconfig_dict.keys()): |
|
|
raise ValueError("qconfig_dict should only contain names(str) as keys.") |
|
|
scripted_qconfig_dict = script_qconfig_dict(qconfig_dict) |
|
|
model = fuse_conv_bn_jit(model, inplace) |
|
|
model_c = torch._C._jit_pass_insert_observers( |
|
|
model._c, "forward", scripted_qconfig_dict, inplace, quant_type |
|
|
) |
|
|
if inplace: |
|
|
model._reconstruct(model_c) |
|
|
else: |
|
|
model = wrap_cpp_module(model_c) |
|
|
return model |
|
|
|
|
|
|
|
|
def _prepare_ondevice_jit( |
|
|
model, |
|
|
qconfig_dict, |
|
|
method_name="forward", |
|
|
inplace=False, |
|
|
quant_type=QuantType.STATIC, |
|
|
): |
|
|
_check_is_script_module(model) |
|
|
if not all(isinstance(x, str) for x in qconfig_dict.keys()): |
|
|
raise ValueError("qconfig_dict should only contain names(str) as keys.") |
|
|
scripted_qconfig_dict = script_qconfig_dict(qconfig_dict) |
|
|
method_graph = model._c._get_method(method_name).graph |
|
|
torch._C._jit_pass_inline(method_graph) |
|
|
model = fuse_conv_bn_jit(model, inplace) |
|
|
model_c = torch._C._jit_pass_insert_observer_method_for_ondevice_ptq( |
|
|
model._c, method_name, scripted_qconfig_dict, inplace, quant_type |
|
|
) |
|
|
if inplace: |
|
|
model._reconstruct(model_c) |
|
|
else: |
|
|
model = wrap_cpp_module(model_c) |
|
|
return model |
|
|
|
|
|
|
|
|
def prepare_jit(model, qconfig_dict, inplace=False): |
|
|
torch._C._log_api_usage_once("quantization_api.quantize_jit.prepare_jit") |
|
|
return _prepare_jit(model, qconfig_dict, inplace, quant_type=QuantType.STATIC) |
|
|
|
|
|
|
|
|
def prepare_dynamic_jit(model, qconfig_dict, inplace=False): |
|
|
torch._C._log_api_usage_once("quantization_api.quantize_jit.prepare_dynamic_jit") |
|
|
return _prepare_jit(model, qconfig_dict, inplace, quant_type=QuantType.DYNAMIC) |
|
|
|
|
|
|
|
|
def _prepare_ondevice_dynamic_jit( |
|
|
model, qconfig_dict, method_name="forward", inplace=False |
|
|
): |
|
|
return _prepare_ondevice_jit( |
|
|
model, qconfig_dict, method_name, inplace, quant_type=QuantType.DYNAMIC |
|
|
) |
|
|
|
|
|
|
|
|
def _convert_jit( |
|
|
model, inplace=False, debug=False, quant_type=QuantType.STATIC, preserved_attrs=None |
|
|
): |
|
|
_check_is_script_module(model) |
|
|
model.eval() |
|
|
model_c = model._c |
|
|
model_c = torch._C._jit_pass_insert_quant_dequant( |
|
|
model_c, "forward", inplace, debug, quant_type |
|
|
) |
|
|
if not debug: |
|
|
is_xpu = all(p.device.type == "xpu" for p in model.parameters()) |
|
|
if not is_xpu: |
|
|
|
|
|
|
|
|
model.cpu() |
|
|
if preserved_attrs is None: |
|
|
preserved_attrs = [] |
|
|
model_c = torch._C._jit_pass_quant_finalize( |
|
|
model_c, quant_type, preserved_attrs |
|
|
) |
|
|
if inplace: |
|
|
model._reconstruct(model_c) |
|
|
else: |
|
|
model = wrap_cpp_module(model_c) |
|
|
torch._C._jit_pass_constant_propagation(model.graph) |
|
|
torch._C._jit_pass_dce(model.graph) |
|
|
return model |
|
|
|
|
|
|
|
|
def _convert_ondevice_jit( |
|
|
model, method_name, inplace=False, debug=False, quant_type=QuantType.STATIC |
|
|
): |
|
|
_check_is_script_module(model) |
|
|
assert quant_type == QuantType.DYNAMIC, ( |
|
|
"This API, while should work for static quant, is only tested for dynamic quant." |
|
|
) |
|
|
assert not method_name.startswith("observe_"), ( |
|
|
"Pass in valid method to be quantized, e.g. forward" |
|
|
) |
|
|
observe_method_name = "observe_" + method_name |
|
|
quantize_method_name = "quantize_" + method_name |
|
|
model_c = model._c |
|
|
model_c = torch._C._jit_pass_insert_quant_dequant_for_ondevice_ptq( |
|
|
model._c, observe_method_name, inplace, debug, QuantType.DYNAMIC |
|
|
) |
|
|
model_c = torch._C._jit_pass_quant_finalize_for_ondevice_ptq( |
|
|
model_c, QuantType.DYNAMIC, quantize_method_name |
|
|
) |
|
|
if inplace: |
|
|
model._reconstruct(model_c) |
|
|
else: |
|
|
model = wrap_cpp_module(model_c) |
|
|
return model |
|
|
|
|
|
|
|
|
def convert_jit(model, inplace=False, debug=False, preserved_attrs=None): |
|
|
torch._C._log_api_usage_once("quantization_api.quantize_jit.convert_jit") |
|
|
return _convert_jit( |
|
|
model, |
|
|
inplace, |
|
|
debug, |
|
|
quant_type=QuantType.STATIC, |
|
|
preserved_attrs=preserved_attrs, |
|
|
) |
|
|
|
|
|
|
|
|
def convert_dynamic_jit(model, inplace=False, debug=False, preserved_attrs=None): |
|
|
torch._C._log_api_usage_once("quantization_api.quantize_jit.convert_dynamic_jit") |
|
|
return _convert_jit( |
|
|
model, |
|
|
inplace, |
|
|
debug, |
|
|
quant_type=QuantType.DYNAMIC, |
|
|
preserved_attrs=preserved_attrs, |
|
|
) |
|
|
|
|
|
|
|
|
def _convert_ondevice_dynamic_jit(model, method_name, inplace=False, debug=False): |
|
|
return _convert_ondevice_jit( |
|
|
model, method_name, inplace, debug, quant_type=QuantType.DYNAMIC |
|
|
) |
|
|
|
|
|
|
|
|
def _quantize_ondevice_dynamic_jit_impl( |
|
|
model, qconfig_dict, method_name, inplace=False |
|
|
): |
|
|
model = _prepare_ondevice_dynamic_jit(model, qconfig_dict, method_name, inplace) |
|
|
model = _convert_ondevice_dynamic_jit(model, method_name, inplace) |
|
|
return model |
|
|
|
|
|
|
|
|
def _quantize_jit( |
|
|
model, |
|
|
qconfig_dict, |
|
|
run_fn=None, |
|
|
run_args=None, |
|
|
inplace=False, |
|
|
debug=False, |
|
|
quant_type=QuantType.STATIC, |
|
|
): |
|
|
|
|
|
|
|
|
if quant_type == QuantType.DYNAMIC: |
|
|
model = prepare_dynamic_jit(model, qconfig_dict, inplace) |
|
|
model = convert_dynamic_jit(model, True, debug) |
|
|
else: |
|
|
assert run_fn, ( |
|
|
"Must provide calibration function for post training static quantization" |
|
|
) |
|
|
assert run_args, ( |
|
|
"Must provide calibration dataset for post training static quantization" |
|
|
) |
|
|
model = prepare_jit(model, qconfig_dict, inplace) |
|
|
run_fn(model, *run_args) |
|
|
model = convert_jit(model, True, debug) |
|
|
|
|
|
torch._C._jit_pass_constant_propagation(model.graph) |
|
|
torch._C._jit_pass_dce(model.graph) |
|
|
return model |
|
|
|
|
|
|
|
|
def quantize_jit(model, qconfig_dict, run_fn, run_args, inplace=False, debug=False): |
|
|
r"""Quantize the input float TorchScript model with |
|
|
post training static quantization. |
|
|
|
|
|
First it will prepare the model for calibration, then it calls |
|
|
`run_fn` which will run the calibration step, after that we will |
|
|
convert the model to a quantized model. |
|
|
|
|
|
Args: |
|
|
`model`: input float TorchScript model |
|
|
`qconfig_dict`: qconfig_dict is a dictionary with names of sub modules as key and |
|
|
qconfig for that module as value, empty key means the qconfig will be applied |
|
|
to whole model unless it's overwritten by more specific configurations, the |
|
|
qconfig for each module is either found in the dictionary or fallback to |
|
|
the qconfig of parent module. |
|
|
|
|
|
Right now qconfig_dict is the only way to configure how the model is quantized, |
|
|
and it is done in the granularity of module, that is, we only support one type |
|
|
of qconfig for each torch.nn.Module, and the qconfig for sub module will |
|
|
override the qconfig for parent module, empty string means global configuration. |
|
|
`run_fn`: a calibration function for calibrating the prepared model |
|
|
`run_args`: positional arguments for `run_fn` |
|
|
`inplace`: carry out model transformations in-place, the original module is |
|
|
mutated |
|
|
`debug`: flag for producing a debug friendly model (preserve weight attribute) |
|
|
|
|
|
Return: |
|
|
Quantized TorchSciprt model. |
|
|
|
|
|
Example: |
|
|
```python |
|
|
import torch |
|
|
from torch.ao.quantization import get_default_qconfig |
|
|
from torch.ao.quantization import quantize_jit |
|
|
|
|
|
ts_model = torch.jit.script( |
|
|
float_model.eval() |
|
|
) # or torch.jit.trace(float_model, input) |
|
|
qconfig = get_default_qconfig("fbgemm") |
|
|
|
|
|
|
|
|
def calibrate(model, data_loader): |
|
|
model.eval() |
|
|
with torch.no_grad(): |
|
|
for image, target in data_loader: |
|
|
model(image) |
|
|
|
|
|
|
|
|
quantized_model = quantize_jit( |
|
|
ts_model, {"": qconfig}, calibrate, [data_loader_test] |
|
|
) |
|
|
``` |
|
|
""" |
|
|
torch._C._log_api_usage_once("quantization_api.quantize_jit.quantize_jit") |
|
|
return _quantize_jit( |
|
|
model, |
|
|
qconfig_dict, |
|
|
run_fn, |
|
|
run_args, |
|
|
inplace, |
|
|
debug, |
|
|
quant_type=QuantType.STATIC, |
|
|
) |
|
|
|
|
|
|
|
|
def quantize_dynamic_jit(model, qconfig_dict, inplace=False, debug=False): |
|
|
r"""Quantize the input float TorchScript model with |
|
|
post training dynamic quantization. |
|
|
Currently only qint8 quantization of torch.nn.Linear is supported. |
|
|
|
|
|
Args: |
|
|
`model`: input float TorchScript model |
|
|
`qconfig_dict`: qconfig_dict is a dictionary with names of sub modules as key and |
|
|
qconfig for that module as value, please see detailed |
|
|
descriptions in :func:`~torch.ao.quantization.quantize_jit` |
|
|
`inplace`: carry out model transformations in-place, the original module is |
|
|
mutated |
|
|
`debug`: flag for producing a debug friendly model (preserve weight attribute) |
|
|
|
|
|
Return: |
|
|
Quantized TorchSciprt model. |
|
|
|
|
|
Example: |
|
|
```python |
|
|
import torch |
|
|
from torch.ao.quantization import per_channel_dynamic_qconfig |
|
|
from torch.ao.quantization import quantize_dynamic_jit |
|
|
|
|
|
ts_model = torch.jit.script( |
|
|
float_model.eval() |
|
|
) # or torch.jit.trace(float_model, input) |
|
|
qconfig = get_default_qconfig("fbgemm") |
|
|
|
|
|
|
|
|
def calibrate(model, data_loader): |
|
|
model.eval() |
|
|
with torch.no_grad(): |
|
|
for image, target in data_loader: |
|
|
model(image) |
|
|
|
|
|
|
|
|
quantized_model = quantize_dynamic_jit( |
|
|
ts_model, {"": qconfig}, calibrate, [data_loader_test] |
|
|
) |
|
|
``` |
|
|
""" |
|
|
torch._C._log_api_usage_once("quantization_api.quantize_jit.quantize_dynamic_jit") |
|
|
return _quantize_jit( |
|
|
model, qconfig_dict, inplace=inplace, debug=debug, quant_type=QuantType.DYNAMIC |
|
|
) |
|
|
|
|
|
|
|
|
def _quantize_ondevice_dynamic_jit( |
|
|
model, qconfig_dict, method_name="forward", inplace=False |
|
|
): |
|
|
r"""Prepares the input float TorchScript model with |
|
|
*on-device* post training dynamic quantization. |
|
|
Currently only qint8 quantization of torch.nn.Linear is supported. |
|
|
|
|
|
Args: |
|
|
`model`: input float TorchScript model |
|
|
`qconfig_dict`: qconfig_dict is a dictionary with names of sub modules as key and |
|
|
qconfig for that module as value, please see detailed |
|
|
`method_name`: Name of the method within the model, to be prepared for quantization |
|
|
descriptions in :func:`~torch.ao.quantization.quantize_jit` |
|
|
`inplace`: carry out model transformations in-place, the original module is |
|
|
mutated |
|
|
|
|
|
Return: |
|
|
TorchScript model that is ready for on device quantization. |
|
|
This means that the returned |
|
|
model has: |
|
|
- Method is inlined. |
|
|
- Model has observer modules inserted in the model. |
|
|
- Model has packed params inserted in the model. However they are empty as in they dont |
|
|
contain valid quantized weights. |
|
|
- observe_<method_name> is added that observe the values to be quantized. |
|
|
- reset_observers_<method_name> to reset observers. |
|
|
- quantize_<method_name> is added to the model. |
|
|
- This method extract scale, zero points. |
|
|
- Quantizes observed weights. |
|
|
- Creates packed params from it and update the attribute of the model with the new values |
|
|
for the packed params. |
|
|
- Reset the original fp32 weights with empty tensor using SetAttr. |
|
|
- quantized_<method_name> is added to the model. |
|
|
- This method uses quantized weights and quantized linear ops instead of fp32 op. |
|
|
- This method should be used for inference post PTQ. |
|
|
- Note that all method's signatures should be the same as method_name. |
|
|
|
|
|
Later on device: |
|
|
- Run reset_observers_<method_name> |
|
|
- Run observe_<method_name> |
|
|
- Run quantize_<method_name> |
|
|
- Now model can be saved and loaded later. |
|
|
- Run model with quantized_<method_name> |
|
|
|
|
|
Example: |
|
|
```python |
|
|
import torch |
|
|
from torch.ao.quantization import per_channel_dynamic_qconfig |
|
|
from torch.ao.quantization.quantize_jit import _quantize_ondevice_dynamic_jit |
|
|
|
|
|
ts_model = torch.jit.script( |
|
|
float_model.eval() |
|
|
) # or torch.jit.trace(float_model, input) |
|
|
qconfig = get_default_qconfig("fbgemm") |
|
|
quant_ready_model = _quantize_ondevice_dynamic_jit( |
|
|
ts_model, {"": qconfig}, "forward", True |
|
|
) |
|
|
``` |
|
|
""" |
|
|
return _quantize_ondevice_dynamic_jit_impl( |
|
|
model, qconfig_dict, method_name, inplace=inplace |
|
|
) |
|
|
|
