PrepareCustomConfig

class torch.ao.quantization.fx.custom_config.PrepareCustomConfig

Custom configuration for prepare_fx() and prepare_qat_fx().

Example usage:

prepare_custom_config = PrepareCustomConfig()             .set_standalone_module_name("module1", qconfig_mapping, example_inputs,                 child_prepare_custom_config, backend_config)             .set_standalone_module_class(MyStandaloneModule, qconfig_mapping, example_inputs,                 child_prepare_custom_config, backend_config)             .set_float_to_observed_mapping(FloatCustomModule, ObservedCustomModule)             .set_non_traceable_module_names(["module2", "module3"])             .set_non_traceable_module_classes([NonTraceableModule1, NonTraceableModule2])             .set_input_quantized_indexes([0])             .set_output_quantized_indexes([0])             .set_preserved_attributes(["attr1", "attr2"])

classmethod from_dict(prepare_custom_config_dict)

Create a PrepareCustomConfig from a dictionary with the following items:

“standalone_module_name”: a list of (module_name, qconfig_mapping, example_inputs, child_prepare_custom_config, backend_config) tuples

“standalone_module_class” a list of (module_class, qconfig_mapping, example_inputs, child_prepare_custom_config, backend_config) tuples

“float_to_observed_custom_module_class”: a nested dictionary mapping from quantization mode to an inner mapping from float module classes to observed module classes, e.g. {“static”: {FloatCustomModule: ObservedCustomModule}}

“non_traceable_module_name”: a list of modules names that are not symbolically traceable “non_traceable_module_class”: a list of module classes that are not symbolically traceable “input_quantized_idxs”: a list of indexes of graph inputs that should be quantized “output_quantized_idxs”: a list of indexes of graph outputs that should be quantized “preserved_attributes”: a list of attributes that persist even if they are not used in forward

This function is primarily for backward compatibility and may be removed in the future.

Return type

PrepareCustomConfig

set_float_to_observed_mapping(float_class, observed_class, quant_type=QuantType.STATIC)

Set the mapping from a custom float module class to a custom observed module class.

The observed module class must have a from_float class method that converts the float module class to the observed module class. This is currently only supported for static quantization.

Return type

PrepareCustomConfig

set_input_quantized_indexes(indexes)

Set the indexes of the inputs of the graph that should be quantized. Inputs are otherwise assumed to be in fp32 by default instead.

Return type

PrepareCustomConfig

set_non_traceable_module_classes(module_classes)

Set the modules that are not symbolically traceable, identified by class.

Return type

PrepareCustomConfig

set_non_traceable_module_names(module_names)

Set the modules that are not symbolically traceable, identified by name.

Return type

PrepareCustomConfig

set_output_quantized_indexes(indexes)

Set the indexes of the outputs of the graph that should be quantized. Outputs are otherwise assumed to be in fp32 by default instead.

Return type

PrepareCustomConfig

set_preserved_attributes(attributes)

Set the names of the attributes that will persist in the graph module even if they are not used in the model’s forward method.

Return type

PrepareCustomConfig

set_standalone_module_class(module_class, qconfig_mapping, example_inputs, prepare_custom_config, backend_config)

Set the configuration for running a standalone module identified by module_class.

If qconfig_mapping is None, the parent qconfig_mapping will be used instead. If prepare_custom_config is None, an empty PrepareCustomConfig will be used. If backend_config is None, the parent backend_config will be used instead.

Return type

PrepareCustomConfig

set_standalone_module_name(module_name, qconfig_mapping, example_inputs, prepare_custom_config, backend_config)

Set the configuration for running a standalone module identified by module_name.

If qconfig_mapping is None, the parent qconfig_mapping will be used instead. If prepare_custom_config is None, an empty PrepareCustomConfig will be used. If backend_config is None, the parent backend_config will be used instead.

Return type

PrepareCustomConfig

to_dict()

Convert this PrepareCustomConfig to a dictionary with the items described in from_dict().

Return type

Dict[str, Any]