Source code for torch.ao.quantization.qconfig_mapping
from__future__importannotationsfromcollectionsimportOrderedDictfromtypingimportAny,Callable,Dict,Tuple,Union,Listimporttorchfrom.fake_quantizeimport(default_weight_fake_quant,FixedQParamsFakeQuantize,)from.observerimport(_PartialWrapper,default_fixed_qparams_range_0to1_observer,default_fixed_qparams_range_neg1to1_observer,default_placeholder_observer,default_weight_observer,)from.qconfigimport(default_reuse_input_qconfig,default_symmetric_qnnpack_qconfig,default_symmetric_qnnpack_qat_qconfig,get_default_qconfig,get_default_qat_qconfig,QConfig,QConfigAny,default_quint8_weight_qconfig)__all__=["get_default_qconfig_mapping","get_default_qat_qconfig_mapping","QConfigMapping",]# TODO: replace all usages with these constants_GLOBAL_DICT_KEY=""_OBJECT_TYPE_DICT_KEY="object_type"_MODULE_NAME_REGEX_DICT_KEY="module_name_regex"_MODULE_NAME_DICT_KEY="module_name"_MODULE_NAME_OBJECT_TYPE_ORDER_DICT_KEY="module_name_object_type_order"# TODO: derive this map from the BackendConfig_FIXED_QPARAMS_OP_TO_OBSERVER:Dict[Union[Callable,str],_PartialWrapper]={torch.nn.Hardsigmoid:default_fixed_qparams_range_0to1_observer,torch.nn.functional.hardsigmoid:default_fixed_qparams_range_0to1_observer,"hardsigmoid":default_fixed_qparams_range_0to1_observer,"hardsigmoid_":default_fixed_qparams_range_0to1_observer,torch.nn.Sigmoid:default_fixed_qparams_range_0to1_observer,torch.sigmoid:default_fixed_qparams_range_0to1_observer,"sigmoid":default_fixed_qparams_range_0to1_observer,"sigmoid_":default_fixed_qparams_range_0to1_observer,torch.nn.Softmax:default_fixed_qparams_range_0to1_observer,torch.nn.Tanh:default_fixed_qparams_range_neg1to1_observer,torch.tanh:default_fixed_qparams_range_neg1to1_observer,"tanh":default_fixed_qparams_range_neg1to1_observer,"tanh_":default_fixed_qparams_range_neg1to1_observer,}def_get_default_qconfig_mapping(is_qat:bool,backend:str,version:int)->QConfigMapping:""" Return the default QConfigMapping for the given quantization type and backend. """ifis_qat:qconfig=get_default_qat_qconfig(backend,version)else:qconfig=get_default_qconfig(backend,version)default_weight=default_weight_fake_quantifis_qatelsedefault_weight_observer# default_per_channel_weight_observer is not currently compatible with fbgemm backend# so we have to modify the weight observer to default_weight_observer or another# per tensor supported observer.# see https://github.com/pytorch/pytorch/issues/47535ifbackendin("fbgemm","x86"):qconfig_transpose=QConfig(activation=qconfig.activation,weight=default_weight)else:qconfig_transpose=qconfig# currently layernorm only supports float weights# we have to add this because otherwise there will be a extra quantize-dequantize pairqconfig_layernorm=QConfig(activation=qconfig.activation,weight=default_placeholder_observer)qconfig_mapping=QConfigMapping() \
.set_global(qconfig) \
.set_object_type("reshape",default_reuse_input_qconfig) \
.set_object_type(torch.nn.ConvTranspose1d,qconfig_transpose) \
.set_object_type(torch.nn.ConvTranspose2d,qconfig_transpose) \
.set_object_type(torch.nn.ConvTranspose3d,qconfig_transpose) \
.set_object_type(torch.nn.functional.conv_transpose1d,qconfig_transpose) \
.set_object_type(torch.nn.functional.conv_transpose2d,qconfig_transpose) \
.set_object_type(torch.nn.functional.conv_transpose3d,qconfig_transpose) \
.set_object_type(torch.nn.functional.layer_norm,qconfig_layernorm) \
.set_object_type(torch.nn.LayerNorm,qconfig_layernorm) \
.set_object_type(torch.nn.PReLU,default_quint8_weight_qconfig) \
# Use special observers for ops with fixed qparamsfixed_qparams_observer_to_qconfig:Dict[Any,QConfigAny]={}forfixed_qparams_op,observerin_FIXED_QPARAMS_OP_TO_OBSERVER.items():ifobserverinfixed_qparams_observer_to_qconfig:fixed_qparams_qconfig=fixed_qparams_observer_to_qconfig[observer]else:ifis_qat:activation=FixedQParamsFakeQuantize.with_args(observer=observer)else:activation=observerfixed_qparams_qconfig=QConfig(activation=activation,weight=default_weight)fixed_qparams_observer_to_qconfig[observer]=fixed_qparams_qconfigqconfig_mapping.set_object_type(fixed_qparams_op,fixed_qparams_qconfig)# TODO Currently it's required that separate ops in a fused op/module have the same qconfig.# Need to be able to support fusion of ops with different qconfigsreturnqconfig_mapping
[docs]defget_default_qconfig_mapping(backend="x86",version=0)->QConfigMapping:""" Return the default QConfigMapping for post training quantization. Args: * ``backend`` (str) : the quantization backend for the default qconfig mapping, should be one of ["x86" (default), "fbgemm", "qnnpack", "onednn"] * ``version`` (int) : the version for the default qconfig mapping """# TODO: add assert for backend choicesreturn_get_default_qconfig_mapping(False,backend,version)
[docs]defget_default_qat_qconfig_mapping(backend="x86",version=1)->QConfigMapping:""" Return the default QConfigMapping for quantization aware training. Args: * ``backend`` (str) : the quantization backend for the default qconfig mapping, should be one of ["x86" (default), "fbgemm", "qnnpack", "onednn"] * ``version`` (int) : the version for the default qconfig mapping """return_get_default_qconfig_mapping(True,backend,version)
def_get_symmetric_qnnpack_qconfig_mapping()->QConfigMapping:""" Return a QConfigMapping that uses `torch.ao.quantization.default_symmetric_qnnpack_qconfig` as the default QConfig. """default_qconfig=default_symmetric_qnnpack_qconfigreturn_get_default_qconfig_mapping_with_default_qconfig(False,"qnnpack",default_qconfig)def_get_symmetric_qnnpack_qat_qconfig_mapping()->QConfigMapping:""" Return a QConfigMapping that uses `torch.ao.quantization.default_symmetric_qnnpack_qat_qconfig` as the default QConfig. """default_qconfig=default_symmetric_qnnpack_qat_qconfigreturn_get_default_qconfig_mapping_with_default_qconfig(True,"qnnpack",default_qconfig)def_get_default_qconfig_mapping_with_default_qconfig(is_qat:bool,backend:str,default_qconfig:QConfig,)->QConfigMapping:""" Return a QConfigMapping that uses the provided qconfig as the default QConfig. """ifis_qat:qconfig_mapping=get_default_qat_qconfig_mapping(backend)else:qconfig_mapping=get_default_qconfig_mapping(backend)qconfig_mapping.set_global(default_qconfig)forpatterninqconfig_mapping.object_type_qconfigs.keys():ifpatternnotin_FIXED_QPARAMS_OP_TO_OBSERVER:qconfig_mapping.set_object_type(pattern,default_qconfig)returnqconfig_mapping_QCONFIG_STYLE_ORDER:List[str]=["global_qconfig","object_type_qconfigs","module_name_regex_qconfigs","module_name_qconfigs","module_name_object_type_order_qconfigs",]
[docs]classQConfigMapping:""" Mapping from model ops to :class:`torch.ao.quantization.QConfig` s. The user can specify QConfigs using the following methods (in increasing match priority): ``set_global`` : sets the global (default) QConfig ``set_object_type`` : sets the QConfig for a given module type, function, or method name ``set_module_name_regex`` : sets the QConfig for modules matching the given regex string ``set_module_name`` : sets the QConfig for modules matching the given module name ``set_module_name_object_type_order`` : sets the QConfig for modules matching a combination of the given module name, object type, and the index at which the module appears Example usage:: qconfig_mapping = QConfigMapping() .set_global(global_qconfig) .set_object_type(torch.nn.Linear, qconfig1) .set_object_type(torch.nn.ReLU, qconfig1) .set_module_name_regex("foo.*bar.*conv[0-9]+", qconfig1) .set_module_name_regex("foo.*", qconfig2) .set_module_name("module1", qconfig1) .set_module_name("module2", qconfig2) .set_module_name_object_type_order("foo.bar", torch.nn.functional.linear, 0, qconfig3) """def__init__(self):# In increasing match priority:self.global_qconfig:QConfigAny=Noneself.object_type_qconfigs:OrderedDict[Union[Callable,str],QConfigAny]=OrderedDict()self.module_name_regex_qconfigs:OrderedDict[str,QConfigAny]=OrderedDict()self.module_name_qconfigs:OrderedDict[str,QConfigAny]=OrderedDict()self.module_name_object_type_order_qconfigs:OrderedDict[Tuple[str,Callable,int],QConfigAny]=\
OrderedDict()
[docs]defset_global(self,global_qconfig:QConfigAny)->QConfigMapping:""" Set the global (default) QConfig. """self.global_qconfig=global_qconfigreturnself
[docs]defset_object_type(self,object_type:Union[Callable,str],qconfig:QConfigAny)->QConfigMapping:""" Set the QConfig for a given module type, function, or method name. If the QConfig for an existing object type was already set, the new QConfig will override the old one. """self.object_type_qconfigs[object_type]=qconfigreturnself
[docs]defset_module_name_regex(self,module_name_regex:str,qconfig:QConfigAny)->QConfigMapping:""" Set the QConfig for modules matching the given regex string. Regexes will be matched in the order in which they are registered through this method. Thus, the caller should register more specific patterns first, e.g.:: qconfig_mapping = QConfigMapping() .set_module_name_regex("foo.*bar.*conv[0-9]+", qconfig1) .set_module_name_regex("foo.*bar.*", qconfig2) .set_module_name_regex("foo.*", qconfig3) In this example, "foo.bar.conv0" would match qconfig1, "foo.bar.linear" would match qconfig2, and "foo.baz.relu" would match qconfig3. If the QConfig for an existing module name regex was already set, the new QConfig will override the old one while preserving the order in which the regexes were originally registered. """self.module_name_regex_qconfigs[module_name_regex]=qconfigreturnself
[docs]defset_module_name(self,module_name:str,qconfig:QConfigAny)->QConfigMapping:""" Set the QConfig for modules matching the given module name. If the QConfig for an existing module name was already set, the new QConfig will override the old one. """self.module_name_qconfigs[module_name]=qconfigreturnself
[docs]defset_module_name_object_type_order(self,module_name:str,object_type:Callable,index:int,qconfig:QConfigAny)->QConfigMapping:""" Set the QConfig for modules matching a combination of the given module name, object type, and the index at which the module appears. If the QConfig for an existing (module name, object type, index) was already set, the new QConfig will override the old one. """self.module_name_object_type_order_qconfigs[(module_name,object_type,index)]=qconfigreturnself
def__repr__(self)->str:output=self.__class__.__name__+" ("forstyle_namein_QCONFIG_STYLE_ORDER:output+=f"\n{style_name}"qconfigs=getattr(self,style_name)ifisinstance(qconfigs,OrderedDict)andlen(qconfigs)>0:forkey,qconfiginqconfigs.items():output+=f"\n{key}: {qconfig}"else:output+=f"\n{qconfigs}"returnoutput+"\n)"# TODO: remove this
[docs]defto_dict(self)->Dict[str,Any]:""" Convert this ``QConfigMapping`` to a dictionary with the following keys: "" (for global QConfig) "object_type" "module_name_regex" "module_name" "module_name_object_type_order" The values of this dictionary are lists of tuples. """return{_GLOBAL_DICT_KEY:self.global_qconfig,_OBJECT_TYPE_DICT_KEY:list(self.object_type_qconfigs.items()),_MODULE_NAME_REGEX_DICT_KEY:list(self.module_name_regex_qconfigs.items()),_MODULE_NAME_DICT_KEY:list(self.module_name_qconfigs.items()),_MODULE_NAME_OBJECT_TYPE_ORDER_DICT_KEY:[(*k,v)fork,vinself.module_name_object_type_order_qconfigs.items()],}
# TODO: remove this
[docs]@classmethoddeffrom_dict(cls,qconfig_dict:Dict[str,Any])->QConfigMapping:""" Create a ``QConfigMapping`` from a dictionary with the following keys (all optional): "" (for global QConfig) "object_type" "module_name_regex" "module_name" "module_name_object_type_order" The values of this dictionary are expected to be lists of tuples. """conf=cls()if_GLOBAL_DICT_KEYinqconfig_dict:conf.set_global(qconfig_dict[_GLOBAL_DICT_KEY])forobject_type,qconfiginqconfig_dict.get(_OBJECT_TYPE_DICT_KEY,[]):conf.set_object_type(object_type,qconfig)formodule_name_regex,qconfiginqconfig_dict.get(_MODULE_NAME_REGEX_DICT_KEY,[]):conf.set_module_name_regex(module_name_regex,qconfig)formodule_name,qconfiginqconfig_dict.get(_MODULE_NAME_DICT_KEY,[]):conf.set_module_name(module_name,qconfig)formodule_name,object_type,index,qconfiginqconfig_dict.get(_MODULE_NAME_OBJECT_TYPE_ORDER_DICT_KEY,[]):conf.set_module_name_object_type_order(module_name,object_type,index,qconfig)returnconf
Docs
Access comprehensive developer documentation for PyTorch
To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.