frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add Custom Module Support List (#82606)

Browse Source 
Summary:
Add a global custon module support list  for the users to specify the modules they want the equalization process support.

To use this list, import it from the _equalize.py file and append module in it.

Unittest passed to check global support list:

https://pxl.cl/28RKG

Test Plan: buck1 test mode/dev //on_device_ai/odai/tests/transforms:test_transforms -- --exact 'on_device_ai/odai/tests/transforms:test_transforms - test_custom_support_list (on_device_ai.odai.tests.transforms.test_input_weight_for_turing.TestInputWeight)'

Reviewed By: jerryzh168

Differential Revision: D38264244

Pull Request resolved: https://github.com/pytorch/pytorch/pull/82606
Approved by: https://github.com/HDCharles
This commit is contained in:
Hao Li
2022-08-03 17:48:51 +00:00
committed by PyTorch MergeBot
parent 9d228fe517
commit aa40503954
2 changed files with 21 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
torch/quantization/fx/_equalize.py
View File

@ -17,6 +17,7 @@ from torch.ao.quantization.fx._equalize import (
default_equalization_qconfig,
fused_module_supports_equalization,
nn_module_supports_equalization,
custom_module_supports_equalization,
node_supports_equalization,
is_equalization_observer,
get_op_node_and_weight_eq_obs,
@ -32,5 +33,6 @@ from torch.ao.quantization.fx._equalize import (
convert_eq_obs,
_convert_equalization_ref,
get_layer_sqnr_dict,
get_equalization_qconfig_dict
get_equalization_qconfig_dict,
CUSTOM_MODULE_SUPP_LIST,
)