pytorch
78144b9f - [Quant][fx][bc-breaking] Replace *custom_config_dict with config objects

Commit
2 years ago
[Quant][fx][bc-breaking] Replace *custom_config_dict with config objects Pull Request resolved: https://github.com/pytorch/pytorch/pull/79066 Following https://github.com/pytorch/pytorch/pull/78452, this commit replaces the following config dicts with python objects: - prepare_custom_config_dict -> PrepareCustomConfig - convert_custom_config_dict -> ConvertCustomConfig - fuse_custom_config_dict -> FuseCustomConfig This leads to better type safety and better user experience in notebook settings due to improved auto completion. The new APIs are as follows: ``` from torch.ao.quantization.fx.custom_config import PrepareCustomConfig prepare_custom_config = PrepareCustomConfig() \ .set_float_to_observed_mapping(float_class, observed_class) \ .set_non_traceable_module_names(["mod1", "mod2"]) \ .set_non_traceable_module_classes([class1, class2]) \ .set_input_quantized_indexes([0, 1]) \ .set_output_quantized_indexes([0]) \ .set_preserved_attributes(["attr1", "attr2"]) convert_custom_config = ConvertCustomConfig() \ .set_observed_to_quantized_mapping(observed_class, quantized_class) \ .set_preserved_attributes(["attr1", "attr2"]) model = prepare_fx( model, qconfig_mapping, example_inputs, prepare_custom_config=prepare_custom_config) model(data) model = convert_fx(model, convert_custom_config=convert_custom_config) ``` For backwards compatibility, prepare_fx, prepare_qat_fx, and convert_fx will continue to accept Dicts, which will be converted to the relevant *CustomConfig object internally. Note that this commit does not modify existing tests to use the new API; they will continue to pass in Dicts as before, which still works but triggers a deprecation warning. This will be handled in a future commit. Differential Revision: [D37088095](https://our.internmc.facebook.com/intern/diff/D37088095/) Approved by: https://github.com/jerryzh168
Author
Committer
Parents
Loading