class_docstr: Implementation of Quantization-Aware Training (QAT) for PyTorch models.
  
  :param model: (Quantizable) Model that needs to be quantized
  :type model: torch.nn.Module
  :param train_loader: DataLoader used for training
  :type train_loader: torch.utils.data.DataLoader
  :param test_loader: DataLoader used for testing
  :type test_loader: torch.utils.data.DataLoader
  :param optimizer: Optimizer for training
  :type optimizer: torch.optim.*
  :param qconfig: Configuration used for quantization
  :type qconfig: Qconfig
  :param criterion: Loss function used for training
  :type criterion: Loss_fn
  :param device: Device used for training ("cpu" or "cuda")
  :type device: torch.device
docstr: Function used for quantization
  
  :param num_train_epochs: Number of epochs used for training
  :type num_train_epochs: int
  :param num_train_batches: Number of batches used for training
  :type num_train_batches: int
  :param param_freeze_epoch: Epoch after which quantizer parameters need to be freezed
  :type param_freeze_epoch: int
  :param bn_freeze_epoch: Epoch after which batch norm mean and variance stats are freezed
  :type bn_freeze_epoch: int
docstr_tok: ["Function","used","for","quantization",":","param","num_train_epochs",":","Number","of","epochs","used","for","training",":","type","num_train_epochs",":","int",":","param","num_train_batches",":","Number","of","batches","used","for","training",":","type","num_train_batches",":","int",":","param","param_freeze_epoch",":","Epoch","after","which","quantizer","parameters","need","to","be","freezed",":","type","param_freeze_epoch",":","int",":","param","bn_freeze_epoch",":","Epoch","after","which","batch","norm","mean","and","variance","stats","are","freezed",":","type","bn_freeze_epoch",":","int"]
code: def quantize(
      self,
      num_train_epochs=10,
      num_train_batches=10,
      param_freeze_epoch=3,
      bn_freeze_epoch=2,
  ):
      """
      Function used for quantization
  
      :param num_train_epochs: Number of epochs used for training
      :type num_train_epochs: int
      :param num_train_batches: Number of batches used for training
      :type num_train_batches: int
      :param param_freeze_epoch: Epoch after which quantizer parameters need to be freezed
      :type param_freeze_epoch: int
      :param bn_freeze_epoch: Epoch after which batch norm mean and variance stats are freezed
      :type bn_freeze_epoch: int
      """
  
      qat_model = deepcopy(self.model)
      qat_model.fuse_model()
  
      optimizer = deepcopy(self.optimizer)
      optimizer.params = qat_model.parameters()
  
      qat_model.qconfig = self.qconfig
  
      torch.quantization.prepare_qat(qat_model, inplace=True)
  
      print("Training model...")
  
      for epoch in range(num_train_epochs):
          print(f"Epoch {epoch}")
          loss, acc = self._train_model(
              qat_model, optimizer, num_train_batches
          )
          print(f"Training Loss: {loss} | Training Acc: {acc}")
  
          if epoch > param_freeze_epoch:
              qat_model.apply(torch.quantization.disable_observer)
  
          if epoch > bn_freeze_epoch:
              qat_model.apply(torch.nn.intrinsic.qat.freeze_bn_stats)
  
          self.quantized_model = torch.quantization.convert(
              qat_model.eval(), inplace=False
          )
          acc = self._evaluate_model(self.quantized_model)
  
          print(f"Evaluation accuracy: {acc}")
  
      return self.quantized_model
  
file_path: KD_Lib/KD_Lib/Quantization/qat/qat.py
imports: from copy import deepcopy
  import torch
  from KD_Lib.Quantization.common import Quantizer
label_desc: Use image node_id 2 for calling the QAT_Quantizer obj's underlying member method code with example usage: obj.quantize(num_train_epochs, num_train_batches, param_freeze_epoch, bn_freeze_epoch) and returns: self