From 7ffcf1550374af29f5b3ce316d4e576423945be0 Mon Sep 17 00:00:00 2001 From: Jerry Zhang Date: Wed, 1 Sep 2021 14:56:14 -0700 Subject: [PATCH] [quant][graphmode][api] Add backend_config_dict to prepare_fx api (#64135) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/64135 We want to start aligning the api with the design in https://github.com/pytorch/pytorch/wiki/Extending-PyTorch-Quantization-to-Custom-Backends We plan to gradually move things from `prepare_custom_config_dict` and `convert_custom_config_dict` to `backend_config_dict` and allow custom backend developer to define their own way of quantizing operators. Test Plan: python test/test_quantization.py TestQuantizeFx python test/test_quantization.py TestQuantizeFxOps Imported from OSS Reviewed By: zou3519 Differential Revision: D30699456 fbshipit-source-id: e3c068da8d3da2270f57719f7159cc71cafa8598 --- torch/quantization/fx/prepare.py | 1 + torch/quantization/quantize_fx.py | 28 +++++++++++++++++++--------- 2 files changed, 20 insertions(+), 9 deletions(-) diff --git a/torch/quantization/fx/prepare.py b/torch/quantization/fx/prepare.py index a6fd660..fb526d0 100644 --- a/torch/quantization/fx/prepare.py +++ b/torch/quantization/fx/prepare.py @@ -1114,6 +1114,7 @@ def prepare( node_name_to_scope: Dict[str, Tuple[str, type]], prepare_custom_config_dict: Optional[Dict[str, Any]] = None, equalization_qconfig_dict: Optional[Dict[str, Any]] = None, + backend_config_dict: Optional[Dict[str, Any]] = None, is_standalone_module: bool = False) -> ObservedGraphModule: """ standalone_module means it a submodule that is not inlined in parent module, and will be quantized separately as one unit. diff --git a/torch/quantization/quantize_fx.py b/torch/quantization/quantize_fx.py index aa8edbb..2dd98ea 100644 --- a/torch/quantization/quantize_fx.py +++ b/torch/quantization/quantize_fx.py @@ -140,8 +140,9 @@ class QuantizationTracer(Tracer): return node def _prepare_fx(model: torch.nn.Module, qconfig_dict: Any, - prepare_custom_config_dict: Dict[str, Any] = None, - equalization_qconfig_dict: Dict[str, Any] = None, + prepare_custom_config_dict: Optional[Dict[str, Any]] = None, + equalization_qconfig_dict: Optional[Dict[str, Any]] = None, + backend_config_dict: Optional[Dict[str, Any]] = None, is_standalone_module: bool = False) -> ObservedGraphModule: r""" Internal helper function for prepare_fx Args: @@ -203,7 +204,8 @@ forward graph of the parent module, def _prepare_standalone_module_fx( model: torch.nn.Module, qconfig_dict: Any, - prepare_custom_config_dict: Dict[str, Any] = None) -> GraphModule: + prepare_custom_config_dict: Dict[str, Any] = None, + backend_config_dict: Dict[str, Any] = None) -> GraphModule: r""" [Internal use only] Prepare a standalone module, so that it can be used when quantizing the parent module. standalone_module means it a submodule that is not inlined in parent module, @@ -224,7 +226,7 @@ def _prepare_standalone_module_fx( same as input_quantized_idxs configuration provided for the standalone module """ - return _prepare_fx(model, qconfig_dict, prepare_custom_config_dict, is_standalone_module=True) + return _prepare_fx(model, qconfig_dict, prepare_custom_config_dict, backend_config_dict, is_standalone_module=True) def fuse_fx(model: torch.nn.Module, fuse_custom_config_dict: Dict[str, Any] = None) -> GraphModule: @@ -265,8 +267,9 @@ def fuse_fx(model: torch.nn.Module, def prepare_fx( model: torch.nn.Module, qconfig_dict: Any, - prepare_custom_config_dict: Dict[str, Any] = None, - equalization_qconfig_dict: Dict[str, Any] = None) -> ObservedGraphModule: + prepare_custom_config_dict: Optional[Dict[str, Any]] = None, + equalization_qconfig_dict: Optional[Dict[str, Any]] = None, + backend_config_dict: Optional[Dict[str, Any]] = None) -> ObservedGraphModule: r""" Prepare a model for post training static quantization Args: @@ -392,6 +395,11 @@ def prepare_fx( with a similar structure as qconfig_dict except it will contain configurations specific to equalization techniques such as input-weight equalization. + `backend_config_dict`: a dictionary that specifies how operators are quantized + in a backend, this includes how the operaetors are observed, + supported fusion patterns, how quantize/dequantize ops are + inserted, supported dtypes etc. The structure of the dictionary is still WIP + and will change in the future, please don't use right now. Return: @@ -420,16 +428,18 @@ def prepare_fx( torch._C._log_api_usage_once("quantization_api.quantize_fx.prepare_fx") assert not model.training, 'prepare_fx only works for models in ' + \ 'eval mode' - return _prepare_fx(model, qconfig_dict, prepare_custom_config_dict, equalization_qconfig_dict) + return _prepare_fx(model, qconfig_dict, prepare_custom_config_dict, equalization_qconfig_dict, backend_config_dict) def prepare_qat_fx( model: torch.nn.Module, qconfig_dict: Any, - prepare_custom_config_dict: Dict[str, Any] = None) -> ObservedGraphModule: + prepare_custom_config_dict: Optional[Dict[str, Any]] = None, + backend_config_dict: Optional[Dict[str, Any]] = None) -> ObservedGraphModule: r""" Prepare a model for quantization aware training Args: `model`: torch.nn.Module model, must be in train mode `qconfig_dict`: see :func:`~torch.quantization.prepare_fx` `prepare_custom_config_dict`: see :func:`~torch.quantization.prepare_fx` + `backend_config_dict`: see :func:`~torch.quantization.prepare_fx` Return: A GraphModule with fake quant modules (configured by qconfig_dict), ready for @@ -457,7 +467,7 @@ def prepare_qat_fx( torch._C._log_api_usage_once("quantization_api.quantize_fx.prepare_qat_fx") assert model.training, 'prepare_qat_fx only works for models in ' + \ 'train mode' - return _prepare_fx(model, qconfig_dict, prepare_custom_config_dict) + return _prepare_fx(model, qconfig_dict, prepare_custom_config_dict, backend_config_dict) def _convert_fx( graph_module: GraphModule, is_reference: bool, -- 2.7.4