'''simple docstring'''





import logging
import re

import pytorch_quantization
import pytorch_quantization.nn as quant_nn
import torch
from pytorch_quantization import calib
from pytorch_quantization.tensor_quant import QuantDescriptor


lowerCAmelCase_	:		Tuple			=     logging.getLogger(__name__)

lowerCAmelCase_	:		Dict			=     50  # max width of layer names
lowerCAmelCase_	:		Union[str, Any]			=     70  # max width of quantizer names
def 						__A							(     UpperCAmelCase					)     ->			Any:




						'''simple docstring'''
						_UpperCamelCase			:			Dict   			=						parser.add_argument_group("quant_trainer arguments"					)
						group.add_argument("--wprec"     ,type=UpperCAmelCase     ,default=8     ,help="weight precision"					)
						group.add_argument("--aprec"     ,type=UpperCAmelCase     ,default=8     ,help="activation precision"					)
						group.add_argument("--quant-per-tensor"     ,action="store_true"     ,help="per tensor weight scaling"					)
						group.add_argument("--quant-disable"     ,action="store_true"     ,help="disable all quantizers"					)
						group.add_argument("--quant-disable-embeddings"     ,action="store_true"     ,help="disable all embeddings quantizers"					)
						group.add_argument("--quant-disable-keyword"     ,type=UpperCAmelCase     ,nargs="+"     ,help="disable quantizers by keyword"					)
						group.add_argument("--quant-disable-layer-module"     ,type=UpperCAmelCase     ,help="disable quantizers by keyword under layer."					)
						group.add_argument("--quant-enable-layer-module"     ,type=UpperCAmelCase     ,help="enable quantizers by keyword under layer"					)
						group.add_argument("--calibrator"     ,default="max"     ,help="which quantization range calibrator to use"					)
						group.add_argument("--percentile"     ,default=UpperCAmelCase     ,type=UpperCAmelCase     ,help="percentile for PercentileCalibrator"					)
						group.add_argument("--fuse-qkv"     ,action="store_true"     ,help="use the same scale factor for qkv"					)
						group.add_argument("--clip-gelu"     ,metavar="N"     ,type=UpperCAmelCase     ,help="clip gelu output maximum value to N"					)
						group.add_argument(
						    "--recalibrate-weights"     ,action="store_true"     ,help=(
						        "recalibrate weight amaxes by taking the max of the weights."
						        " amaxes will be computed with the current quantization granularity (axis)."
						    )     ,)
def 						__A							(     UpperCAmelCase					)     ->			Optional[Any]:




						'''simple docstring'''

						if args.calibrator == "max":
												_UpperCamelCase			:			List[Any]   			=						"max"
						elif args.calibrator == "percentile":
												if args.percentile is None:
																		raise ValueError("Specify --percentile when using percentile calibrator"					)
												_UpperCamelCase			:			Tuple   			=						"histogram"
						elif args.calibrator == "mse":
												_UpperCamelCase			:			int   			=						"histogram"
						else:
												raise ValueError(f'''Invalid calibrator {args.calibrator}'''					)

						_UpperCamelCase			:			int   			=						QuantDescriptor(num_bits=args.aprec     ,calib_method=UpperCAmelCase					)
						_UpperCamelCase			:			Any   			=						QuantDescriptor(num_bits=args.wprec     ,axis=(None if args.quant_per_tensor else (0,))					)
						quant_nn.QuantLinear.set_default_quant_desc_input(UpperCAmelCase					)
						quant_nn.QuantLinear.set_default_quant_desc_weight(UpperCAmelCase					)
def 						__A							(     UpperCAmelCase     ,UpperCAmelCase     ,UpperCAmelCase=False     ,UpperCAmelCase=False					)     ->			List[Any]:




						'''simple docstring'''
						logger.info("Configuring Model for Quantization"					)
						logger.info(f'''using quantization package {pytorch_quantization.__file__}'''					)

						if not calib:
												if args.quant_disable_embeddings:
																		set_quantizer_by_name(UpperCAmelCase     ,["embeddings"]     ,which="weight"     ,_disabled=UpperCAmelCase					)

												if args.quant_disable:
																		set_quantizer_by_name(UpperCAmelCase     ,[""]     ,_disabled=UpperCAmelCase					)

												if args.quant_disable_keyword:
																		set_quantizer_by_name(UpperCAmelCase     ,args.quant_disable_keyword     ,_disabled=UpperCAmelCase					)

												if args.quant_disable_layer_module:
																		set_quantizer_by_name(UpperCAmelCase     ,[R"layer.\d+." + args.quant_disable_layer_module]     ,_disabled=UpperCAmelCase					)

												if args.quant_enable_layer_module:
																		set_quantizer_by_name(UpperCAmelCase     ,[R"layer.\d+." + args.quant_enable_layer_module]     ,_disabled=UpperCAmelCase					)

												if args.recalibrate_weights:
																		recalibrate_weights(UpperCAmelCase					)

												if args.fuse_qkv:
																		fuse_qkv(UpperCAmelCase     ,UpperCAmelCase					)

						if args.clip_gelu:
												clip_gelu(UpperCAmelCase     ,args.clip_gelu					)

						# if args.local_rank in [-1, 0] and not calib:
						print_quant_summary(UpperCAmelCase					)
def 						__A							(     UpperCAmelCase					)     ->			Optional[int]:




						'''simple docstring'''
						logger.info("Enabling Calibration"					)
						for name, module in model.named_modules():
												if name.endswith("_quantizer"					):
																		if module._calibrator is not None:
																								module.disable_quant()
																								module.enable_calib()
																		else:
																								module.disable()
																		logger.info(f'''{name:80}: {module}'''					)
def 						__A							(     UpperCAmelCase     ,UpperCAmelCase					)     ->			Any:




						'''simple docstring'''
						logger.info("Loading calibrated amax"					)
						for name, module in model.named_modules():
												if name.endswith("_quantizer"					):
																		if module._calibrator is not None:
																								if isinstance(module._calibrator     ,calib.MaxCalibrator					):
																														module.load_calib_amax()
																								else:
																														module.load_calib_amax("percentile"     ,percentile=args.percentile					)
																								module.enable_quant()
																								module.disable_calib()
																		else:
																								module.enable()
						model.cuda()
						print_quant_summary(UpperCAmelCase					)
def 						__A							(     UpperCAmelCase     ,UpperCAmelCase					)     ->			Optional[Any]:




						'''simple docstring'''

						def fusea(UpperCAmelCase     ,UpperCAmelCase     ,UpperCAmelCase					):
												for mod in [qq, qk, qv]:
																		if not hasattr(UpperCAmelCase     ,"_amax"					):
																								print("          WARNING: NO AMAX BUFFER"					)
																								return
												_UpperCamelCase			:			Tuple   			=						qq._amax.detach().item()
												_UpperCamelCase			:			Optional[Any]   			=						qk._amax.detach().item()
												_UpperCamelCase			:			Optional[int]   			=						qv._amax.detach().item()

												_UpperCamelCase			:			Union[str, Any]   			=						max(UpperCAmelCase     ,UpperCAmelCase     ,UpperCAmelCase					)
												qq._amax.fill_(UpperCAmelCase					)
												qk._amax.fill_(UpperCAmelCase					)
												qv._amax.fill_(UpperCAmelCase					)
												logger.info(f'''          q={q:5.2f} k={k:5.2f} v={v:5.2f} -> {amax:5.2f}'''					)

						for name, mod in model.named_modules():
												if name.endswith(".attention.self"					):
																		logger.info(f'''FUSE_QKV: {name:{name_width}}'''					)
																		fusea(mod.matmul_q_input_quantizer     ,mod.matmul_k_input_quantizer     ,mod.matmul_v_input_quantizer					)
																		if args.quant_per_tensor:
																								fusea(mod.query._weight_quantizer     ,mod.key._weight_quantizer     ,mod.value._weight_quantizer					)
def 						__A							(     UpperCAmelCase     ,UpperCAmelCase					)     ->			Any:




						'''simple docstring'''

						for name, mod in model.named_modules():
												if name.endswith(".output.dense"					) and not name.endswith("attention.output.dense"					):
																		_UpperCamelCase			:			Dict   			=						mod._input_quantizer._amax.data.detach().item()
																		mod._input_quantizer._amax.data.detach().clamp_(max=UpperCAmelCase					)
																		_UpperCamelCase			:			Union[str, Any]   			=						mod._input_quantizer._amax.data.detach().item()
																		logger.info(f'''CLIP_GELU: {name:{name_width}} amax: {amax_init:5.2f} -> {amax:5.2f}'''					)
def 						__A							(     UpperCAmelCase					)     ->			Union[str, Any]:




						'''simple docstring'''

						for name, mod in model.named_modules():
												if hasattr(UpperCAmelCase     ,"_weight_quantizer"					) and mod._weight_quantizer.axis is not None:
																		_UpperCamelCase			:			Optional[int]   			=						mod.weight.shape[0]
																		_UpperCamelCase			:			Tuple   			=						mod._weight_quantizer._amax.detach()
																		_UpperCamelCase			:			Union[str, Any]   			=						torch.ones(UpperCAmelCase     ,dtype=amax.dtype     ,device=amax.device					) * amax
																		print(f'''expanding {name} {amax} -> {mod._weight_quantizer._amax}'''					)
def 						__A							(     UpperCAmelCase					)     ->			Optional[Any]:




						'''simple docstring'''

						for name, mod in model.named_modules():
												if hasattr(UpperCAmelCase     ,"_weight_quantizer"					):
																		if not hasattr(mod.weight_quantizer     ,"_amax"					):
																								print("RECALIB: {name:{name_width}} WARNING: NO AMAX BUFFER"					)
																								continue

																		# determine which axes to reduce across
																		# e.g. a 4D tensor quantized per axis 0 should reduce over (1,2,3)
																		_UpperCamelCase			:			Union[str, Any]   			=						set() if mod._weight_quantizer.axis is None else set(mod._weight_quantizer.axis					)
																		_UpperCamelCase			:			str   			=						set(range(len(mod.weight.size()					)					)					) - axis_set
																		_UpperCamelCase			:			int   			=						pytorch_quantization.utils.reduce_amax(mod.weight     ,axis=UpperCAmelCase     ,keepdims=UpperCAmelCase					).detach()
																		logger.info(f'''RECALIB: {name:{name_width}} {mod._weight_quantizer._amax.flatten()} -> {amax.flatten()}'''					)
																		_UpperCamelCase			:			Any   			=						amax
def 						__A							(     UpperCAmelCase     ,UpperCAmelCase=2_5     ,UpperCAmelCase=1_8_0     ,UpperCAmelCase=None					)     ->			Union[str, Any]:




						'''simple docstring'''

						if ignore is None:
												_UpperCamelCase			:			int   			=						[]
						elif not isinstance(UpperCAmelCase     ,UpperCAmelCase					):
												_UpperCamelCase			:			Tuple   			=						[ignore]

						_UpperCamelCase			:			str   			=						0
						for name, mod in model.named_modules():
												if not hasattr(UpperCAmelCase     ,"weight"					):
																		continue
												_UpperCamelCase			:			Any   			=						max(UpperCAmelCase     ,len(UpperCAmelCase					)					)

						for name, mod in model.named_modules():
												_UpperCamelCase			:			str   			=						getattr(UpperCAmelCase     ,"_input_quantizer"     ,UpperCAmelCase					)
												_UpperCamelCase			:			Union[str, Any]   			=						getattr(UpperCAmelCase     ,"_weight_quantizer"     ,UpperCAmelCase					)
												if not hasattr(UpperCAmelCase     ,"weight"					):
																		continue
												if type(UpperCAmelCase					) in ignore:
																		continue
												if [True for s in ignore if type(UpperCAmelCase					) is str and s in name]:
																		continue
												_UpperCamelCase			:			Tuple   			=						f'''Act:{input_q.extra_repr()}'''
												_UpperCamelCase			:			Any   			=						f'''Wgt:{weight_q.extra_repr()}'''
												_UpperCamelCase			:			Tuple   			=						f'''{name:{name_width}} {act_str} {wgt_str}'''
												if len(UpperCAmelCase					) <= line_width:
																		logger.info(UpperCAmelCase					)
												else:
																		logger.info(f'''{name:{name_width}} {act_str}'''					)
																		logger.info(f'''{'  ':{name_width}} {wgt_str}'''					)
def 						__A							(     UpperCAmelCase					)     ->			List[str]:




						'''simple docstring'''
						_UpperCamelCase			:			Dict   			=						0
						for name, mod in model.named_modules():
												if isinstance(UpperCAmelCase     ,pytorch_quantization.nn.TensorQuantizer					):
																		print(f'''{name:80} {mod}'''					)
																		count += 1
						print(f'''{count} TensorQuantizers found in model'''					)
def 						__A							(     UpperCAmelCase     ,UpperCAmelCase     ,UpperCAmelCase     ,UpperCAmelCase     ,UpperCAmelCase					)     ->			Union[str, Any]:




						'''simple docstring'''
						_UpperCamelCase			:			int   			=						getattr(UpperCAmelCase     ,UpperCAmelCase     ,UpperCAmelCase					)
						if quantizer_mod is not None:
												assert hasattr(UpperCAmelCase     ,UpperCAmelCase					)
												setattr(UpperCAmelCase     ,UpperCAmelCase     ,UpperCAmelCase					)
						else:
												logger.warning(f'''{name} has no {quantizer}'''					)
def 						__A							(     UpperCAmelCase     ,UpperCAmelCase     ,UpperCAmelCase="both"     ,**UpperCAmelCase					)     ->			Optional[Any]:




						'''simple docstring'''
						_UpperCamelCase			:			List[str]   			=						f'''Warning: changing {which} quantizers of {name:{qname_width}}'''
						for k, v in kwargs.items():
												s += f''' {k}={v}'''
												if which in ["input", "both"]:
																		set_quantizer(UpperCAmelCase     ,UpperCAmelCase     ,"_input_quantizer"     ,UpperCAmelCase     ,UpperCAmelCase					)
												if which in ["weight", "both"]:
																		set_quantizer(UpperCAmelCase     ,UpperCAmelCase     ,"_weight_quantizer"     ,UpperCAmelCase     ,UpperCAmelCase					)
						logger.info(UpperCAmelCase					)






def 						__A							(     UpperCAmelCase     ,UpperCAmelCase     ,**UpperCAmelCase					)     ->			Dict:




						'''simple docstring'''

						for name, mod in model.named_modules():
												if hasattr(UpperCAmelCase     ,"_input_quantizer"					) or hasattr(UpperCAmelCase     ,"_weight_quantizer"					):
																		for n in names:
																								if re.search(UpperCAmelCase     ,UpperCAmelCase					):
																														set_quantizers(UpperCAmelCase     ,UpperCAmelCase     ,**UpperCAmelCase					)
												elif name.endswith("_quantizer"					):
																		for n in names:
																								if re.search(UpperCAmelCase     ,UpperCAmelCase					):
																														_UpperCamelCase			:			Optional[int]   			=						f'''Warning: changing {name:{name_width}}'''
																														for k, v in kwargs.items():
																																				s += f''' {k}={v}'''
																																				setattr(UpperCAmelCase     ,UpperCAmelCase     ,UpperCAmelCase					)
																														logger.info(UpperCAmelCase					)



'''simple docstring'''





from collections import OrderedDict
from typing import List, Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


lowerCAmelCase_	:		Optional[Any]			=     logging.get_logger(__name__)

lowerCAmelCase_	:		List[Any]			=     {
    """google/efficientnet-b7""": """https://huggingface.co/google/efficientnet-b7/resolve/main/config.json""",
}




class 		SCREAMING_SNAKE_CASE     (					SCREAMING_SNAKE_CASE	):







				'''simple docstring'''
				UpperCAmelCase__							=							'''efficientnet'''





				def __init__(						self					:						Optional[int]       ,		lowercase__					:						int = 3       ,		lowercase__					:						int = 600       ,		lowercase__					:						float = 2.0       ,		lowercase__					:						float = 3.1       ,		lowercase__					:						int = 8       ,		lowercase__					:						List[int] = [3, 3, 5, 3, 5, 5, 3]       ,		lowercase__					:						List[int] = [32, 16, 24, 40, 80, 112, 192]       ,		lowercase__					:						List[int] = [16, 24, 40, 80, 112, 192, 320]       ,		lowercase__					:						List[int] = []       ,		lowercase__					:						List[int] = [1, 2, 2, 2, 1, 2, 1]       ,		lowercase__					:						List[int] = [1, 2, 2, 3, 3, 4, 1]       ,		lowercase__					:						List[int] = [1, 6, 6, 6, 6, 6, 6]       ,		lowercase__					:						float = 0.2_5       ,		lowercase__					:						str = "swish"       ,		lowercase__					:						int = 2_560       ,		lowercase__					:						str = "mean"       ,		lowercase__					:						float = 0.0_2       ,		lowercase__					:						float = 0.0_0_1       ,		lowercase__					:						float = 0.9_9       ,		lowercase__					:						float = 0.5       ,		lowercase__					:						float = 0.2       ,		**lowercase__					:						int       ,		)							->Optional[int]:


										'''simple docstring'''




										super().__init__(**lowercase__					)

										_UpperCamelCase			:			int   			=						num_channels
										_UpperCamelCase			:			Any   			=						image_size
										_UpperCamelCase			:			Union[str, Any]   			=						width_coefficient
										_UpperCamelCase			:			List[str]   			=						depth_coefficient
										_UpperCamelCase			:			List[Any]   			=						depth_divisor
										_UpperCamelCase			:			Optional[Any]   			=						kernel_sizes
										_UpperCamelCase			:			int   			=						in_channels
										_UpperCamelCase			:			Any   			=						out_channels
										_UpperCamelCase			:			Union[str, Any]   			=						depthwise_padding
										_UpperCamelCase			:			Optional[Any]   			=						strides
										_UpperCamelCase			:			Union[str, Any]   			=						num_block_repeats
										_UpperCamelCase			:			List[str]   			=						expand_ratios
										_UpperCamelCase			:			Dict   			=						squeeze_expansion_ratio
										_UpperCamelCase			:			Union[str, Any]   			=						hidden_act
										_UpperCamelCase			:			List[str]   			=						hidden_dim
										_UpperCamelCase			:			Dict   			=						pooling_type
										_UpperCamelCase			:			List[str]   			=						initializer_range
										_UpperCamelCase			:			Any   			=						batch_norm_eps
										_UpperCamelCase			:			List[str]   			=						batch_norm_momentum
										_UpperCamelCase			:			Optional[Any]   			=						dropout_rate
										_UpperCamelCase			:			Dict   			=						drop_connect_rate
										_UpperCamelCase			:			Dict   			=						sum(lowercase__					) * 4





class 		SCREAMING_SNAKE_CASE     (					SCREAMING_SNAKE_CASE	):







				'''simple docstring'''
				UpperCAmelCase__							=							version.parse('''1.11'''	)





				@property
				def  snake_case__				(						self					:						Union[str, Any]					)							->Mapping[str, Mapping[int, str]]:


										'''simple docstring'''




										return OrderedDict(
										    [
										        ("pixel_values", {0: "batch", 1: "num_channels", 2: "height", 3: "width"}),
										    ]					)







				@property
				def  snake_case__				(						self					:						Any					)							->float:


										'''simple docstring'''




										return 1e-5


