KD_Lib.Quantization.common.base_class module

class KD_Lib.Quantization.common.base_class.Quantizer(model, qconfig, train_loader=None, test_loader=None, optimizer=None, criterion=None, device=device(type='cpu'))[source]

Bases: object

Basic Implementation of Quantization for PyTorch models.

  • model (torch.nn.Module) – Model that needs to be pruned
  • qconfig (Qconfig) – Configuration used for quantization
  • train_loader (torch.utils.data.DataLoader) – DataLoader used for training
  • test_loader (torch.utils.data.DataLoader) – DataLoader used for testing
  • optimizer (torch.optim.*) – Optimizer for training
  • criterion (Loss_fn) – Loss function used for calibration
  • device (torch.device) – Device used for training (“cpu” or “cuda”)

Function for printing sizes of the original and quantized model


Function used for reporting inference performance of original and quantized models Note that performance here referes to the following: 1. Accuracy achieved on the testset 2. Time taken for evaluating on the testset


Function used for quantization