frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-11-05 16:44:58 +08:00
Code Issues Packages Projects Releases Wiki Activity

[qunat][graphmode][fx] Standalone module takes float as input and output (#48671)

Browse Source 
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/48671

Standalone module might be called separately so it's better to use float
as interface.

Test Plan: Imported from OSS

Reviewed By: vkuzo

Differential Revision: D25256184

fbshipit-source-id: e209492a180ce1f81f31c8d6057956a74bad20b1
This commit is contained in:
Jerry Zhang
2020-12-02 20:33:07 -08:00
committed by Facebook GitHub Bot
parent 22c3ae8b57
commit ea573ea944
5 changed files with 39 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
torch/quantization/quantize_fx.py
View File

@ -99,14 +99,8 @@ def _prepare_standalone_module_fx(model, qconfig_dict, prepare_custom_config_dic
standalone_module means it a submodule that is not inlined in parent module,
and will be quantized separately as one unit.
input of the module is quantized in parent module, output of the module
is quantized in the standalone module.
Extra attributes in output GraphModule while preparing a standalone module:
_standalone_module_observed_input_idxs(List[Int]): a list of indexs for the graph inputs that
needs to be observed in parent module
_output_is_observed(Bool): a boolean variable indicate whether the output of the
custom module is observed or not
Both input and output of the module are observed in the
standalone module.
"""
return _prepare_fx(model, qconfig_dict, prepare_custom_config_dict, is_standalone_module=True)
@ -356,11 +350,8 @@ def _convert_standalone_module_fx(graph_module, debug=False, convert_custom_conf
r""" [Internal use only] Convert a model produced by :func:`~torch.quantization.prepare_standalone_module_fx`
and convert it to a quantized model
The inputs will be quantized by parent module, checks `_standalone_module_observed_input_idxs` of
input model and will treat these inputs as quantized
also will not dequantize the final output
Return:
A quantized standalone module which accepts quantized input(if needed)
and produces quantized output (if needed).
A quantized standalone module which accepts float input
and produces float output.
"""
return _convert_fx(graph_module, debug, convert_custom_config_dict, is_standalone_module=True)