+ "ݳj6^RIt^RIHtHt^RIHt^RIHtHt^RI H t H t H t ^RI HtHtHtHt^RIHt^RIHt^R IHt^R IHt.ROtR R ltR Rlt]P4P6P8P:]P4P6P8P<]P4P6P>P:]P4P@PB.t"RRlt#Rt$RRRllt%R#)N) GraphModuleNode) PassManager)_fold_conv_bn_qat_fuse_conv_bn_qat)DuplicateDQPassPortNodeMetaForQDQ Quantizer)_disallow_eval_train_fuse_conv_bn__fuse_linear_bn__get_node_name_to_scope) constant_fold)#_convert_to_reference_decomposed_fx)prepare) reference_representation_rewritec<V^8dQhR\R\R\/#model quantizerreturnrr )formats"h/home/luca/screen-checking/.venv/lib/python3.14/site-packages/torchao/quantization/pt2e/quantize_pt2e.py __annotate__r$s+gg gggc\PPR4VPp\ V4p\ V4\ V4VPV4pVPV4VPV4\VVRVPR7pVPPV4\V4pVPPFpVP R8XgKVP"\P$P&P(JgKJVP*^,p\-V\.4'gQhVP R8XgQh\-VP"\04'gQhVP3VP"4p\5Wa4p\7WP"V4K VPPFpVP R8XgKVP"\P$P&P8JgKJVP*^,p\-V\.4'gQhVP R8XgQh\-VP"\04'gQhVP3VP"4p \5W4p \7WP"V 4K V#)aPrepare a model for post training quantization Args: * `model` (torch.fx.GraphModule): a model captured by `torch.export.export` API. * `quantizer`: A backend specific quantizer that conveys how user want the model to be quantized. Tutorial for how to write a quantizer can be found here: https://pytorch.org/tutorials/prototype/pt2e_quantizer.html Return: A GraphModule with observer (based on quantizer annotation), ready for calibration Example:: import torch from torchao.quantization.pt2e.quantize_pt2e import prepare_pt2e from torchao.quantization.pt2e.quantizer import ( XNNPACKQuantizer, get_symmetric_quantization_config, ) class M(torch.nn.Module): def __init__(self) -> None: super().__init__() self.linear = torch.nn.Linear(5, 10) def forward(self, x): return self.linear(x) # initialize a floating point model float_model = M().eval() # define calibration function def calibrate(model, data_loader): model.eval() with torch.no_grad(): for image, target in data_loader: model(image) # Step 1. program capture # NOTE: this API will be updated to torch.export API in the future, but the captured # result shoud mostly stay the same m = torch.export.export(m, *example_inputs).module() # we get a model with aten ops # Step 2. quantization # backend developer will write their own Quantizer and expose methods to allow # users to express how they # want the model to be quantized quantizer = XNNPACKQuantizer().set_global(get_symmetric_quantization_config()) m = prepare_pt2e(m, quantizer) # run calibration # calibrate(m, sample_inference_data) z&torchao.quantization.pt2e.prepare_pt2eFis_qatobs_or_fq_callback call_functionget_attr)torch_C_log_api_usage_oncemetarr r transform_for_annotationannotatevalidaterprepare_obs_or_fq_callbackupdater graphnodesoptargetops higher_orderscanargs isinstancerstr get_submodule prepare_pt2esetattr while_loop) rroriginal_graph_metanode_name_to_scopenodescan_combine_fn_nodescan_combine_fnprepared_scan_combine_fnwhile_loop_body_fn_nodewhile_loop_body_fnprepared_while_loop_body_fns && rr8r8$s0t HH  !IJ**075U  . .u 5E u u  $??  E  JJ)*  'E  !! 77o %$++9O9O9T9T*T#'99Q< 2D99 99'**j8 882993?? ??#112F2M2MNO'3O'O $ E668P Q" !! GG & uyy55@@@&*iil #5t<< <<*--; ;;5< None: super().__init__() self.linear = torch.nn.Linear(5, 10) def forward(self, x): return self.linear(x) # initialize a floating point model float_model = M().eval() # define the training loop for quantization aware training def train_loop(model, train_data): model.train() for image, target in data_loader: ... # Step 1. program capture # NOTE: this API will be updated to torch.export API in the future, but the captured # result shoud mostly stay the same m = torch.export.export(m, *example_inputs).module() # we get a model with aten ops # Step 2. quantization # backend developer will write their own Quantizer and expose methods to allow # users to express how they # want the model to be quantized quantizer = XNNPACKQuantizer().set_global(get_symmetric_quantization_config()) m = prepare_qat_pt2e(m, quantizer) # run quantization aware training train_loop(prepared_model, train_loop) z*torchao.quantization.pt2e.prepare_qat_pt2eTr) r$r%r&r'rr(r)r*rrr+r,r )rrr;r<s&& rprepare_qat_pt2erFsp HH  !MN**07  . .u 5E u ue  $??  E  JJ)*  'E Lrc0V^8dQhR\R\/#)rnr)rbool)rs"rrrs>>d>t>rcVVPR8H;'dVP\9#)a@If there is any pure ops between get_attr and quantize op they will be const propagated e.g. get_attr(weight) -> transpose -> quantize -> dequantize* (Note: dequantize op is not going to be constant propagated) This filter is added because we don't want to constant fold the things that are not related to quantization r")r/r0 _QUANT_OPS)rHs&r_quant_node_constraintrLs$ 44? " = =qxx:'==rc^RIHp^RIHp^RIHpRpRpVP 4Fwrg\ V\PPPP4'g}\ V\PPPP4'g?\ V\PPPP4'dRp\ Wq4'g%\ Ws4'g\ Wr4'gKRpK V'dV'dQR4hV'dV'dQR4hV#)r) FakeQuantize)AffineQuantizedObserverBase) ObserverBaseFTz2Cannot be prepared using both torch.ao and torchao) 'torchao.quantization.pt2e.fake_quantizerN"torchao.quantization.pt2e.observerrOrP named_modulesr5r$ao quantization fake_quantizeobserver)rtorchao_FakeQuantize#torchao_AffineQuantizedObserverBasetorchao_ObserverBaseis_torch_ao_preparedis_torchao_prepared_ms& r1_is_torchao_prepared_do_not_use_outside_this_filer_sX ##% q%((//==JJ K K!UXX22;;HHII!UXX22;;WWXX#' q / /!22!AA"& && @ &' @ ' rcHV^8dQhR\R\R\R\/#)rruse_reference_representation fold_quantizer)rrI)rs"rrrs6SS S"&SS Src\PPR4\V\4'g\ RV R24hVP pVPPFpVPR8XgKVP\PPPJgKJVP^,p\V\4'gQhVPR8XgQh\VP\ 4'gQhVP#VP4p\%VVVR7p\'WPV4K VPPFpVPR8XgKVP\PPP(JgKJVP^,p\V\4'gQhVPR8XgQh\VP\ 4'gQhVP#VP4p \%V VVR7p \'WPV 4K \+V4p\-V4p\/\14.4p V !V4P2p\/\54.4p V !V4P2pV'd\7V\84V'd \;V4pVP P=V4\?V4pV#)aConvert a calibrated/trained model to a quantized model Args: * `model` (torch.fx.GraphModule): calibrated/trained model * `use_reference_representation` (bool): boolean flag to indicate whether to produce referece representation or not * `fold_quantize` (bool): boolean flag for whether fold the quantize op or not Returns: quantized model, either in q/dq representation or reference representation Example:: # prepared_model: the model produced by `prepare_pt2e`/`prepare_qat_pt2e` and calibration/training # `convert_pt2e` produces a quantized model that represents quantized computation with # quantize dequantize ops and fp32 ops by default. # Please refer to # https://pytorch.org/tutorials/prototype/pt2e_quant_ptq_static.html#convert-the-calibrated-model-to-a-quantized-model # for detailed explanation of output quantized model quantized_model = convert_pt2e(prepared_model) z&torchao.quantization.pt2e.convert_pt2ezjUnexpected argument type for `use_reference_representation`, please make sure you intend to pass argument z to convert_pt2er"r#)rarb) r$r%r&r5rI ValueErrorr'r-r.r/r0r1r2r3r4rr6r7 convert_pt2er9r:rrrr graph_moduler rrLrr,r ) rrarbr;r=r>r?converted_scan_combine_fnrArBconverted_while_loop_body_fnpms &&& rreres|4 HH  !IJ 2D 9 9 <rys &:T )8N gTJ\ II""66>> II""66== II""77?? II%%  >"JSSr