Spaces:
Running
Running
| 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: | |
| # Moving model parameters to CPU since quantized operators | |
| # are only supported on CPU and XPU right now | |
| 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): | |
| # Always do inplace convert because the Tensor is already | |
| # copied in prepare_jit when inplace is False | |
| 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) | |