class_docstr: Perform quantization aware training (QAT) of models. This algorithm simulates the effects of
  quantization during training, by quantizing and dequantizing the weights and/or activations during
  the model's forward pass. The forward and backward pass computations are conducted in ``float`` dtype,
  however, these ``float`` values follow the constraints imposed by ``int8`` and ``quint8`` dtypes. Thus,
  this algorithm adjusts the model's weights while closely simulating the numerics which get
  executed during quantized inference, allowing model's weights to adjust to quantization
  constraints.
  
  For more details, please refer to  `Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only
  Inference <https://arxiv.org/pdf/1712.05877.pdf>`_.
  
  Example:
  
          .. code-block:: python
  
              import torch.nn as nn
              from coremltools.optimize.torch.quantization import (
                  LinearQuantizer,
                  LinearQuantizerConfig,
              )
  
              model = nn.Sequential(
                  OrderedDict(
                      {
                          "conv": nn.Conv2d(1, 20, (3, 3)),
                          "relu1": nn.ReLU(),
                          "conv2": nn.Conv2d(20, 20, (3, 3)),
                          "relu2": nn.ReLU(),
                      }
                  )
              )
  
              loss_fn = define_loss()
  
              # initialize the quantizer
              config = LinearQuantizerConfig.from_dict(
                  {
                      "global_config": {
                          "quantization_scheme": "symmetric",
                          "milestones": [0, 100, 400, 400],
                      }
                  }
              )
  
              quantizer = LinearQuantizer(model, config)
  
              # prepare the model to insert FakeQuantize layers for QAT
              model = quantizer.prepare()
  
              # use quantizer in your PyTorch training loop
              for inputs, labels in data:
                  output = model(inputs)
                  loss = loss_fn(output, labels)
                  loss.backward()
                  optimizer.step()
                  quantizer.step()
  
              # convert operations to their quanitzed counterparts using parameters learnt via QAT
              model = quantizer.finalize(inplace=True)
  
  Args:
      model (:obj:`torch.nn.Module`): Module to be trained.
      config (:py:class:`_LinearQuantizerConfig`): Config that specifies how
          different submodules in the model will be quantized.
          Default config is used when passed as ``None``.
class_docstr_tok: ["Perform","quantization","aware","training","(","QAT",")","of","models",".","This","algorithm","simulates","the","effects","of","quantization","during","training",",","by","quantizing","and","dequantizing","the","weights","and\/or","activations","during","the","model","'s","forward","pass",".","The","forward","and","backward","pass","computations","are","conducted","in","``","float","``","dtype",",","however",",","these","``","float","``","values","follow","the","constraints","imposed","by","``","int8","``","and","``","quint8","``","dtypes",".","Thus",",","this","algorithm","adjusts","the","model","'s","weights","while","closely","simulating","the","numerics","which","get","executed","during","quantized","inference",",","allowing","model","'s","weights","to","adjust","to","quantization","constraints",".","For","more","details",",","please","refer","to","`","Quantization","and","Training","of","Neural","Networks","for","Efficient","Integer-Arithmetic-Only","Inference","<","https",":","\/\/arxiv.org\/pdf\/1712.05877.pdf",">","`","_",".","Example",":","..","code-block",":",":","python","import","torch.nn","as","nn","from","coremltools.optimize.torch.quantization","import","(","LinearQuantizer",",","LinearQuantizerConfig",",",")","model","=","nn.Sequential","(","OrderedDict","(","{","``","conv","''",":","nn.Conv2d","(","1",",","20",",","(","3",",","3",")",")",",","``","relu1","''",":","nn.ReLU","(",")",",","``","conv2","''",":","nn.Conv2d","(","20",",","20",",","(","3",",","3",")",")",",","``","relu2","''",":","nn.ReLU","(",")",",","}",")",")","loss_fn","=","define_loss","(",")","#","initialize","the","quantizer","config","=","LinearQuantizerConfig.from_dict","(","{","``","global_config","''",":","{","``","quantization_scheme","''",":","``","symmetric","''",",","``","milestones","''",":","[","0",",","100",",","400",",","400","]",",","}","}",")","quantizer","=","LinearQuantizer","(","model",",","config",")","#","prepare","the","model","to","insert","FakeQuantize","layers","for","QAT","model","=","quantizer.prepare","(",")","#","use","quantizer","in","your","PyTorch","training","loop","for","inputs",",","labels","in","data",":","output","=","model","(","inputs",")","loss","=","loss_fn","(","output",",","labels",")","loss.backward","(",")","optimizer.step","(",")","quantizer.step","(",")","#","convert","operations","to","their","quanitzed","counterparts","using","parameters","learnt","via","QAT","model","=","quantizer.finalize","(","inplace=True",")","Args",":","model","(",":","obj",":","`","torch.nn.Module","`",")",":","Module","to","be","trained",".","config","(",":","py",":","class",":","`","_LinearQuantizerConfig","`",")",":","Config","that","specifies","how","different","submodules","in","the","model","will","be","quantized",".","Default","config","is","used","when","passed","as","``","None","``","."]
docstr: Prepares the model for export.
  
  Args:
      model (:py:class:`_torch.nn.Module`): Model to be finalized.
      inplace (:obj:`bool`): If ``True``, model transformations are carried out in-place and
          the original module is mutated; otherwise, a copy of the model is mutated and returned.
  
  .. note::
      Once the model is finalized with ``in_place = True``, it may not be
      runnable on the GPU.
docstr_tok: ["Prepares","the","model","for","export",".","Args",":","model","(",":","py",":","class",":","`","_torch.nn.Module","`",")",":","Model","to","be","finalized",".","inplace","(",":","obj",":","`","bool","`",")",":","If","``","True","``",",","model","transformations","are","carried","out","in-place","and","the","original","module","is","mutated",";","otherwise",",","a","copy","of","the","model","is","mutated","and","returned",".","..","note",":",":","Once","the","model","is","finalized","with","``","in_place","=","True","``",",","it","may","not","be","runnable","on","the","GPU","."]
code: def finalize(
      self,
      model: _Optional[_torch.nn.Module] = None,
      inplace: bool = False,
  ) -> _torch.nn.Module:
      """
      Prepares the model for export.
  
      Args:
          model (:py:class:`_torch.nn.Module`): Model to be finalized.
          inplace (:obj:`bool`): If ``True``, model transformations are carried out in-place and
              the original module is mutated; otherwise, a copy of the model is mutated and returned.
  
      .. note::
          Once the model is finalized with ``in_place = True``, it may not be
          runnable on the GPU.
      """
      if not self._is_prepared:
          _logger.warning(
              "Model has not been prepared for QAT. This API call "
              "will be a no-op. prepare method must be called before "
              "a call to the finalize method."
          )
          return self._model
      if model is None:
          model = self._model
      if not inplace:
          model = _copy.deepcopy(model)
      model.eval()
      finalized_model = _convert_to_reference_fx(
          model,
          qconfig_mapping=self._qconfig_mapping,
          backend_config=_get_backend_config(),
      )
      if model is None:
          self._model = finalized_model
      return finalized_model
  
imports: import copy
  import logging
  from typing import Any
  from typing import Optional
  from typing import Tuple
  import torch
  import torch.ao.quantization
  import torch.nn.intrinsic.qat
  from torch.ao.quantization.fx.custom_config import PrepareCustomConfig
  from torch.ao.quantization.quantize_fx import convert_to_reference_fx
  from coremltools.optimize.torch._utils.math_utils import rmse_error
  from coremltools.optimize.torch._utils.torch_utils import get_eval_model
  from coremltools.optimize.torch.base_model_optimizer import BaseModelOptimizer
  from coremltools.optimize.torch.base_model_optimizer import _Report
  from coremltools.optimize.torch.quantization._backend_config import get_backend_config
  from coremltools.optimize.torch.quantization._backend_config import get_supported_modules
  from coremltools.optimize.torch.quantization._configure import QATConfigurationHandler
  from coremltools.optimize.torch.quantization._qconfig_mapping import _QConfigMappingBuilder
  from coremltools.optimize.torch.quantization.quantization_config import LinearQuantizerConfig
  from coremltools.optimize.torch.quantization.quantization_config import ModuleLinearQuantizerConfig
label_desc: Use image node_id 5 for calling the LinearQuantizer obj's underlying member method code with example usage: obj.finalize(model, inplace) and returns: finalized_model, self