@@ -72,7 +72,7 @@ jobs:
             python3 -m venv venv
             . venv/bin/activate
             python3 -m pip install pip==22.0.4
-            cat requirements.txt | cut -f1 -d"#" | xargs -n 1 -L 1 pip install --progress-bar off
+            cat requirements.txt | cut -f1 -d"#" | grep "^[^--;]" | xargs -n 1 -L 1 pip install --progress-bar off --extra-index-url https://pypi.ngc.nvidia.com
       - run:
           name: edit package version
           command: |

@@ -30,4 +30,6 @@ packaging>=20.4
 # not directly required, pinned by Snyk to avoid a vulnerability
 wheel>=0.38.0
 # not directly required, pinned by Snyk to avoid a vulnerability
-pygments>=2.7.4 
+pygments>=2.7.4
+--extra-index-url https://pypi.ngc.nvidia.com
+pytorch-quantization==2.1.2

@@ -21,7 +21,8 @@ def readme():
 
 def get_requirements():
     with open(REQ_LOCATION, encoding="utf-8") as f:
-        return f.read().splitlines()
+        requirements = f.read().splitlines()
+        return [r for r in requirements if not r.startswith("--") and not r.startswith("#")]
 
 
 def get_pro_requirements():
@@ -45,6 +46,7 @@ setup(
     install_requires=get_requirements(),
     packages=find_packages(where="./src"),
     package_dir={"": "src"},
+    dependency_links=["https://pypi.ngc.nvidia.com"],
     package_data={
         "super_gradients.recipes": ["*.yaml", "**/*.yaml"],
         "super_gradients.common": ["auto_logging/auto_logging_conf.json"],

import torch
from torch import nn

from super_gradients.training.utils.quantization.selective_quantization_utils import SelectiveQuantizer


def non_default_calibrators_example():
    class MyModel(nn.Module):
        def __init__(self) -> None:
            super().__init__()
            self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)

        def forward(self, x):
            return self.conv1(x)

    module = MyModel()

    # Initialize the quantization utility, with different calibrators, and quantize the module
    q_util = SelectiveQuantizer(default_quant_modules_calib_method="percentile", default_per_channel_quant_modules=False)
    q_util.quantize_module(module)

    print(module)  # You should expect to see QuantConv2d, with Histogram calibrators

    x = torch.rand(1, 3, 32, 32)
    with torch.no_grad():
        y = module(x)
        torch.testing.assert_close(y.size(), (1, 8, 32, 32))


if __name__ == "__main__":
    non_default_calibrators_example()

import torch
from pytorch_quantization import nn as quant_nn
from torch import nn

from super_gradients.training.dataloaders import cifar10_train
from super_gradients.training.utils.quantization.calibrator import QuantizationCalibrator
from super_gradients.training.utils.quantization.core import SGQuantMixin
from super_gradients.training.utils.quantization.export import export_quantized_module_to_onnx
from super_gradients.training.utils.quantization.selective_quantization_utils import SelectiveQuantizer


def e2e_example():
    class MyBlock(nn.Module):
        def __init__(self, in_feats, out_feats) -> None:
            super().__init__()
            self.in_feats = in_feats
            self.out_feats = out_feats
            self.flatten = nn.Flatten()
            self.linear = nn.Linear(in_feats, out_feats)

        def forward(self, x):
            return self.linear(self.flatten(x))

    class MyQuantizedBlock(SGQuantMixin):
        def __init__(self, in_feats, out_feats) -> None:
            super().__init__()
            self.flatten = nn.Flatten()
            self.linear = quant_nn.QuantLinear(in_feats, out_feats)

        def forward(self, x):
            return self.linear(self.flatten(x))

    class MyModel(nn.Module):
        def __init__(self, res, n_classes) -> None:
            super().__init__()
            self.my_block = MyBlock(3 * (res**2), n_classes)

        def forward(self, x):
            return self.my_block(x)

    res = 32
    n_clss = 10
    module = MyModel(res, n_clss)

    # QUANTIZE
    q_util = SelectiveQuantizer()
    q_util.register_quantization_mapping(layer_names={"my_block"}, quantized_target_class=MyQuantizedBlock)
    q_util.quantize_module(module)

    # CALIBRATE (PTQ)
    train_loader = cifar10_train()
    calib = QuantizationCalibrator()
    calib.calibrate_model(module, method=q_util.default_quant_modules_calib_method, calib_data_loader=train_loader)

    module.cuda()
    # SANITY
    x = torch.rand(1, 3, res, res, device="cuda")
    with torch.no_grad():
        y = module(x)
        torch.testing.assert_close(y.size(), (1, n_clss))

    print(module)

    # EXPORT TO ONNX
    export_quantized_module_to_onnx(module, "my_quantized_model.onnx", input_shape=(1, 3, res, res))


if __name__ == "__main__":
    e2e_example()

import torch
from pytorch_quantization import nn as quant_nn
from torch import nn

from super_gradients.training.utils.quantization.core import SGQuantMixin
from super_gradients.training.utils.quantization.selective_quantization_utils import SelectiveQuantizer, register_quantized_module


def register_quantization_mapping_with_decorator_example():
    # ARRANGE
    class MyBlock(nn.Module):
        def __init__(self, in_feats, out_feats) -> None:
            super().__init__()
            self.in_feats = in_feats
            self.out_feats = out_feats
            self.flatten = nn.Flatten()
            self.linear = nn.Linear(in_feats, out_feats)

        def forward(self, x):
            return self.linear(self.flatten(x))

    @register_quantized_module(float_source=MyBlock)
    class MyQuantizedBlock(SGQuantMixin):
        def __init__(self, in_feats, out_feats) -> None:
            super().__init__()
            self.flatten = nn.Flatten()
            self.linear = quant_nn.QuantLinear(in_feats, out_feats)

        def forward(self, x):
            return self.linear(self.flatten(x))

    class MyModel(nn.Module):
        def __init__(self, res, n_classes) -> None:
            super().__init__()
            self.conv = nn.Conv2d(3, 4, kernel_size=3, padding=1)
            self.my_block = MyBlock(4 * (res**2), n_classes)

        def forward(self, x):
            y = self.conv(x)
            return self.my_block(y)

    res = 32
    n_clss = 10
    module = MyModel(res, n_clss)

    # TEST
    q_util = SelectiveQuantizer()
    q_util.quantize_module(module)

    x = torch.rand(1, 3, res, res)

    print(module)

    # ASSERT
    with torch.no_grad():
        y = module(x)
        torch.testing.assert_close(y.size(), (1, n_clss))


if __name__ == "__main__":
    register_quantization_mapping_with_decorator_example()

import argparse

from torch import nn

import super_gradients
from super_gradients import Trainer
from super_gradients.training import MultiGPUMode
from super_gradients.training import models as sg_models
from super_gradients.training.dataloaders import imagenet_train, imagenet_val
from super_gradients.training.metrics import Accuracy, Top5
from super_gradients.training.metrics.metric_utils import get_metrics_dict
from super_gradients.training.models.classification_models.resnet import Bottleneck as sg_Bottleneck
from super_gradients.training.utils.quantization.calibrator import QuantizationCalibrator
from super_gradients.training.utils.quantization.core import QuantizedMetadata
from super_gradients.training.utils.quantization.export import export_quantized_module_to_onnx
from super_gradients.modules.quantization.resnet_bottleneck import QuantBottleneck as sg_QuantizedBottleneck
from super_gradients.training.utils.quantization.selective_quantization_utils import SelectiveQuantizer


def naive_quantize(model: nn.Module):
    q_util = SelectiveQuantizer(default_quant_modules_calib_method="max", default_per_channel_quant_modules=True)
    q_util.quantize_module(model)

    return model


def selective_quantize(model: nn.Module):
    mappings = {
        sg_Bottleneck: QuantizedMetadata(
            float_source=sg_Bottleneck,
            quantized_target_class=sg_QuantizedBottleneck,
            action=QuantizedMetadata.ReplacementAction.RECURE_AND_REPLACE,
        ),
    }

    sq_util = SelectiveQuantizer(custom_mappings=mappings, default_quant_modules_calib_method="max", default_per_channel_quant_modules=True)
    sq_util.quantize_module(model)

    return model


def sg_vanilla_resnet50():
    return sg_models.get("resnet50", pretrained_weights="imagenet", num_classes=1000)


def sg_naive_qdq_resnet50():
    return naive_quantize(sg_vanilla_resnet50())


def sg_selective_qdq_resnet50():
    return selective_quantize(sg_vanilla_resnet50())


models = {
    "sg_vanilla_resnet50": sg_vanilla_resnet50,
    "sg_naive_qdq_resnet50": sg_naive_qdq_resnet50,
    "sg_selective_qdq_resnet50": sg_selective_qdq_resnet50,
}

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    super_gradients.init_trainer()

    parser.add_argument("--max_epochs", type=int, default=10)
    parser.add_argument("--lr", type=float, default=0.001)
    parser.add_argument("--batch", type=int, default=128)
    parser.add_argument("--model_name", type=str)
    parser.add_argument("--calibrate", action="store_true")

    args, _ = parser.parse_known_args()

    train_params = {
        "max_epochs": args.max_epochs,
        "initial_lr": args.lr,
        "optimizer": "SGD",
        "optimizer_params": {"weight_decay": 0.0001, "momentum": 0.9, "nesterov": True},
        "loss": "cross_entropy",
        "train_metrics_list": [Accuracy(), Top5()],
        "valid_metrics_list": [Accuracy(), Top5()],
        "test_metrics_list": [Accuracy(), Top5()],
        "loss_logging_items_names": ["Loss"],
        "metric_to_watch": "Accuracy",
        "greater_metric_to_watch_is_better": True,
    }

    trainer = Trainer(experiment_name=args.model_name, multi_gpu=MultiGPUMode.OFF, device="cuda")

    train_dataloader = imagenet_train(dataloader_params={"batch_size": args.batch, "shuffle": True})
    val_dataloader = imagenet_val(dataloader_params={"batch_size": args.batch, "shuffle": True, "drop_last": True})

    model = models[args.model_name]().cuda()

    if args.calibrate:
        calibrator = QuantizationCalibrator(verbose=False)
        calibrator.calibrate_model(model, method="max", calib_data_loader=train_dataloader, num_calib_batches=1024 // args.batch or 1)

    trainer.train(model=model, training_params=train_params, train_loader=train_dataloader, valid_loader=val_dataloader)

    val_results_tuple = trainer.test(model=model, test_loader=val_dataloader, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)
    valid_metrics_dict = get_metrics_dict(val_results_tuple, trainer.test_metrics, trainer.loss_logging_items_names)

    export_quantized_module_to_onnx(model=model, onnx_filename=f"{args.model_name}.onnx", input_shape=(args.batch, 3, 224, 224))

    print(f"FINAL ACCURACY: {valid_metrics_dict['Accuracy'].cpu().item()}")

import torch
from torch import nn

from super_gradients.training.utils.quantization.core import SkipQuantization
from super_gradients.training.utils.quantization.selective_quantization_utils import SelectiveQuantizer


def skipping_quantization_example():
    class MyModel(nn.Module):
        def __init__(self) -> None:
            super().__init__()
            self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)
            self.conv2 = SkipQuantization(nn.Conv2d(8, 8, kernel_size=3, padding=1))  # can use the wrapper to skip

        def forward(self, x):
            return self.conv2(self.conv1(x))

    module = MyModel()

    # Initialize the quantization utility, register layers to skip, and quantize the module
    q_util = SelectiveQuantizer()
    q_util.register_skip_quantization(layer_names={"conv1"})  # can also configure skip by layer names
    q_util.quantize_module(module)

    print(module)  # You should expect to see Conv2d

    x = torch.rand(1, 3, 32, 32)
    with torch.no_grad():
        y = module(x)
        torch.testing.assert_close(y.size(), (1, 8, 32, 32))


if __name__ == "__main__":
    skipping_quantization_example()

import torch
from torch import nn

from super_gradients.training.utils.quantization.selective_quantization_utils import SelectiveQuantizer


def vanilla_quantize_all_example():
    class MyModel(nn.Module):
        def __init__(self) -> None:
            super().__init__()
            self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)

        def forward(self, x):
            return self.conv1(x)

    module = MyModel()

    # Initialize the quantization utility, and quantize the module
    q_util = SelectiveQuantizer()
    q_util.quantize_module(module)

    print(module)  # You should expect to see QuantConv2d

    x = torch.rand(1, 3, 32, 32)
    with torch.no_grad():
        y = module(x)
        torch.testing.assert_close(y.size(), (1, 8, 32, 32))


if __name__ == "__main__":
    vanilla_quantize_all_example()

"""
QAT example for Resnet18

The purpose of this example is to demonstrate the usage of QAT in super_gradients.

Behind the scenes, when passing enable_qat=True, a callback for QAT will be added.

Once triggered, the following will happen:
- The model will be rebuilt with quantized nn.modules.
- The pretrained imagenet weights will be loaded to it.
- We perform calibration with 2 batches from our training set (1024 samples = 8 gpus X 128 samples_per_batch).
- We evaluate the calibrated model (accuracy is logged under calibrated_model_accuracy).
- The calibrated checkpoint prior to QAT is saved under ckpt_calibrated_{calibration_method}.pth.
- We fine tune the calibrated model for 1 epoch.

Finally, once training is over- we trigger a pos-training callback that will export the ONNX files.

"""
from super_gradients.training import Trainer, MultiGPUMode, models, dataloaders
from super_gradients.training.metrics.classification_metrics import Accuracy

import super_gradients
from super_gradients.training.utils.quantization_utils import PostQATConversionCallback

super_gradients.init_trainer()

trainer = Trainer("resnet18_qat_example",
                  multi_gpu=MultiGPUMode.DISTRIBUTED_DATA_PARALLEL)

train_loader = dataloaders.imagenet_train()
valid_loader = dataloaders.imagenet_val()

model = models.get("resnet18", pretrained_weights="imagenet")

train_params = {"max_epochs": 1,
                "lr_mode": "step",
                "optimizer": "SGD",
                "lr_updates": [],
                "lr_decay_factor": 0.1,
                "initial_lr": 0.001, "loss": "cross_entropy",
                "train_metrics_list": [Accuracy()],
                "valid_metrics_list": [Accuracy()],

                "metric_to_watch": "Accuracy",
                "greater_metric_to_watch_is_better": True,
                "average_best_models": False,
                "enable_qat": True,
                "qat_params": {
                    "start_epoch": 0,  # first epoch for quantization aware training.
                    "quant_modules_calib_method": "percentile",
                    # statistics method for amax computation (one of [percentile, mse, entropy, max]).
                    "calibrate": True,  # whether to perform calibration.
                    "num_calib_batches": 2,  # number of batches to collect the statistics from.
                    "percentile": 99.99  # percentile value to use when Trainer,
                },
                "phase_callbacks": [PostQATConversionCallback(dummy_input_size=(1, 3, 224, 224))]
                }

trainer.train(model=model, training_params=train_params, train_loader=train_loader, valid_loader=valid_loader)

from .resnet_bottleneck import QuantBottleneck

__all__ = ["QuantBottleneck"]

from super_gradients.training.models import Bottleneck

try:
    from pytorch_quantization import nn as quant_nn
    from super_gradients.training.utils.quantization.core import SGQuantMixin, QuantizedMetadata
    from super_gradients.training.utils.quantization.selective_quantization_utils import register_quantized_module

    _imported_pytorch_quantization_failure = None
except (ImportError, NameError, ModuleNotFoundError) as import_err:
    _imported_pytorch_quantization_failure = import_err


@register_quantized_module(float_source=Bottleneck, action=QuantizedMetadata.ReplacementAction.RECURE_AND_REPLACE)
class QuantBottleneck(SGQuantMixin):
    """
    we just insert quantized tensor to the shortcut (=residual) layer, so that it would be quantized
    NOTE: we must quantize the float instance, so the mode should be
          QuantizedMetadata.ReplacementAction.RECURE_AND_REPLACE
    """

    if _imported_pytorch_quantization_failure is not None:
        raise _imported_pytorch_quantization_failure

    @classmethod
    def from_float(cls, float_instance: Bottleneck, **kwargs):
        float_instance.shortcut.add_module("residual_quantizer", quant_nn.TensorQuantizer(kwargs.get("quant_desc_input")))
        return float_instance

@@ -36,7 +36,6 @@ from super_gradients.training import utils as core_utils, models, dataloaders
 from super_gradients.training.models import SgModule
 from super_gradients.training.pretrained_models import PRETRAINED_NUM_CLASSES
 from super_gradients.training.utils import sg_trainer_utils
-from super_gradients.training.utils.quantization_utils import QATCallback
 from super_gradients.training.utils.sg_trainer_utils import MonitoredValue, parse_args, log_main_training_params
 from super_gradients.training.exceptions.sg_trainer_exceptions import UnsupportedOptimizerFormat, GPUModeNotSetupError
 from super_gradients.training.losses import LOSSES
@@ -1045,10 +1044,9 @@ class Trainer:
         # ADD CALLBACK FOR QAT
         self.enable_qat = core_utils.get_param(self.training_params, "enable_qat", False)
         if self.enable_qat:
-            self.qat_params = core_utils.get_param(self.training_params, "qat_params")
-            if self.qat_params is None:
-                raise ValueError("Must pass QAT params when enable_qat=True")
-            self.phase_callbacks.append(QATCallback(**self.qat_params))
+            raise NotImplementedError(
+                "QAT is not implemented as a plug-and-play feature yet. Please refer to examples/resnet_qat to learn how to do it manually."
+            )
 
         self.phase_callback_handler = CallbackHandler(callbacks=self.phase_callbacks)
 

from super_gradients.common.abstractions.abstract_logger import get_logger

logger = get_logger(__name__)
try:
    from super_gradients.training.utils.quantization.core import _inject_class_methods_to_default_quant_types

    _inject_class_methods_to_default_quant_types()
except (ImportError, NameError, ModuleNotFoundError):
    logger.warning("Failed to import pytorch_quantization")

"""
Quantization utilities

Methods are based on:
https://github.com/NVIDIA/TensorRT/blob/51a4297753d3e12d0eed864be52400f429a6a94c/tools/pytorch-quantization/examples/torchvision/classification_flow.py#L385

(Licensed under the Apache License, Version 2.0)
"""

import torch
from tqdm import tqdm

from super_gradients.common.abstractions.abstract_logger import get_logger
from super_gradients.training.utils.distributed_training_utils import get_local_rank, get_world_size
from torch.distributed import all_gather

logger = get_logger(__name__)

try:
    from pytorch_quantization import nn as quant_nn
    from pytorch_quantization import calib

    _imported_pytorch_quantization_failure = None
except (ImportError, NameError, ModuleNotFoundError) as import_err:
    logger.warning("Failed to import pytorch_quantization")
    _imported_pytorch_quantization_failure = import_err


class QuantizationCalibrator:
    def __init__(self, verbose: bool = True) -> None:
        if _imported_pytorch_quantization_failure is not None:
            raise _imported_pytorch_quantization_failure
        super().__init__()
        self.verbose = verbose

    def calibrate_model(
        self,
        model: torch.nn.Module,
        calib_data_loader: torch.utils.data.DataLoader,
        method: str = "percentile",
        num_calib_batches: int = 2,
        percentile: float = 99.99,
    ):
        """
        Calibrates torch model with quantized modules.

        :param model:               torch.nn.Module, model to perfrom the calibration on.
        :param calib_data_loader:   torch.utils.data.DataLoader, data loader of the calibration dataset.
                                    Assumes that the first element of the tuple is the input image.
        :param method:              str, One of [percentile, mse, entropy, max].
                                    Statistics method for amax computation of the quantized modules
                                    (Default=percentile).
        :param num_calib_batches:   int, number of batches to collect the statistics from.
        :param percentile:          float, percentile value to use when SgModel,quant_modules_calib_method='percentile'.
                                    Discarded when other methods are used (Default=99.99).

        """
        acceptable_methods = ["percentile", "mse", "entropy", "max"]
        if method in acceptable_methods:
            with torch.no_grad():
                self._collect_stats(model, calib_data_loader, num_batches=num_calib_batches)
                # FOR PERCENTILE WE MUST PASS PERCENTILE VALUE THROUGH KWARGS,
                # SO IT WOULD BE PASSED TO module.load_calib_amax(**kwargs), AND IN OTHER METHODS WE MUST NOT PASS IT.
                if method == "precentile":
                    self._compute_amax(model, method="percentile", percentile=percentile)
                else:
                    self._compute_amax(model, method=method)
        else:
            raise ValueError(f"Unsupported quantization calibration method, " f"expected one of: {'.'.join(acceptable_methods)}, however, received: {method}")

    def _collect_stats(self, model, data_loader, num_batches):
        """Feed data to the network and collect statistics"""
        local_rank = get_local_rank()
        world_size = get_world_size()

        device = next(model.parameters()).device

        # Enable calibrators
        self._enable_calibrators(model)

        # Feed data to the network for collecting stats
        for i, (image, *_) in tqdm(enumerate(data_loader), total=num_batches, disable=local_rank > 0):
            if world_size > 1:
                all_batches = [torch.zeros_like(image, device=device) for _ in range(world_size)]
                all_gather(all_batches, image.to(device=device))
            else:
                all_batches = [image]

            for local_image in all_batches:
                model(local_image.to(device=device))
            if i >= num_batches:
                break

        # Disable calibrators
        self._disable_calibrators(model)

    def _disable_calibrators(self, model):
        for name, module in model.named_modules():
            if isinstance(module, quant_nn.TensorQuantizer):
                if module._calibrator is not None:
                    module.enable_quant()
                    module.disable_calib()
                else:
                    module.enable()

    def _enable_calibrators(self, model):
        for name, module in model.named_modules():
            if isinstance(module, quant_nn.TensorQuantizer):
                if module._calibrator is not None:
                    module.disable_quant()
                    module.enable_calib()
                else:
                    module.disable()

    def _compute_amax(self, model, **kwargs):
        for name, module in model.named_modules():
            if isinstance(module, quant_nn.TensorQuantizer):
                if module._calibrator is not None:
                    if isinstance(module._calibrator, calib.MaxCalibrator):
                        module.load_calib_amax()
                    else:
                        module.load_calib_amax(**kwargs)
                if self.verbose:
                    print(f"{name:40}: {module}")

import inspect
from dataclasses import dataclass
from enum import Enum
from typing import Union, Type, Optional, Set

from pytorch_quantization.nn.modules._utils import QuantMixin, QuantInputMixin
from pytorch_quantization.tensor_quant import QuantDescriptor
from torch import nn


def _extract_init_args(cls, float_instance, ignore_init_args: Set[str] = ()):
    """
    Inspecting the __init__ args, and searching for corresponding properties from the float instance
    e.g., for `__init__(self, a)` the mechanism will look for `float_instance.a` and pass that value to `__init__`
    """
    required_init_params = list(inspect.signature(cls.__init__).parameters)[1:]  # [0] is self

    if "kwargs" in required_init_params:  # we don't want to search for a state named `kwargs`
        required_init_params.pop(required_init_params.index("kwargs"))

    float_instance_state = {}
    for p in required_init_params:
        if p in ignore_init_args:  # ignore these args and don't pick state from the instance
            continue
        if not hasattr(float_instance, p):
            raise ValueError(
                f"{float_instance.__class__.__name__} is missing `{p}` which is required "
                f"in {cls.__name__}.__init__. Either override `SGQuantBase.from_float` "
                f"or add {p} as state for {float_instance.__class__.__name__}."
            )
        float_instance_state[p] = getattr(float_instance, p)

    # Edge-cases
    if "bias" in float_instance_state:
        if float_instance_state["bias"] is None:  # None is the state when bias=False in torch.nn
            float_instance_state["bias"] = False
        elif not isinstance(float_instance_state["bias"], bool):  # Tensor is the state when bias=True in torch.nn
            float_instance_state["bias"] = True
        # in case bias is a boolean - we don't do anything, so it is taken as-is, either True or False
    return float_instance_state


def _from_float(cls, float_instance, ignore_init_args: Set[str] = (), **kwargs):
    init_params = _extract_init_args(cls, float_instance, ignore_init_args)
    init_params.update(**kwargs)
    return cls(**init_params)


class SGQuantMixin(nn.Module):
    """
    A base class for user custom Quantized classes.
    Every Quantized class must inherit this mixin, which adds `from_float` class-method.
    NOTES:
        * the Quantized class may also inherit from the native `QuantMixin` or `QuantInputMixin`
        * quant descriptors (for inputs and weights) will be passed as `kwargs`. The module may ignore them if they are
          not necessary
        * the default implementation of `from_float` is inspecting the __init__ args, and searching for corresponding
          properties from the float instance that is passed as argument, e.g., for `__init__(self, a)`
          the mechanism will look for `float_instance.a` and pass that value to the `__init__` method
    """

    @classmethod
    def from_float(cls, float_instance, **kwargs):
        required_init_params = list(inspect.signature(cls.__init__).parameters)[1:]  # [0] is self

        # if cls.__init__ has explicit `quant_desc_input` or `quant_desc_weight` - we don't search the state of the
        # float module, because it would not contain this state. these values are injected by the framework
        ignore_init_args = {"quant_desc_input", "quant_desc_weight"}.intersection(set(required_init_params))

        # if cls.__init__ doesn't have neither **kwargs, nor `quant_desc_input` and `quant_desc_weight`,
        # we should also remove these keys from the passed kwargs and make sure there's nothing more!
        if "kwargs" not in required_init_params:
            for arg in ("quant_desc_input", "quant_desc_weight"):
                if arg in ignore_init_args:
                    continue
                kwargs.pop(arg, None)  # we ignore if not existing

        return _from_float(cls, float_instance, ignore_init_args, **kwargs)


class SkipQuantization(nn.Module):
    """
    This class wraps a float module instance, and defines that this instance will not be converted to quantized version

    Example:
        self.my_block = SkipQuantization(MyBlock(4, n_classes))
    """

    def __init__(self, module: nn.Module) -> None:
        super().__init__()
        self.float_module = module
        self.forward = module.forward


@dataclass(init=True)
class QuantizedMetadata:
    """
    This dataclass is responsible for holding the information regarding float->quantized module relation.
    It can be both layer-grained and module-grained, e.g.,
    `module.backbone.conv1 -> QuantConv2d`, `nn.Linear -> QuantLinear`, etc...

    Args:
        float_source:               the name of a specific layer (e.g., `module.backbone.conv1`),
                                    or a specific type (e.g., `Conv2d`) that will be later quantized
        quantized_target_class:     the quantized type that the source will be converted to
        action:                     how to resolve the conversion, we either:
                                    - SKIP: skip it,
                                    - UNWRAP: unwrap the instance and work with the wrapped one
                                      (i.e., we wrap with a mapper),
                                    - REPLACE: replace source with an instance of the
                                      quantized type
                                    - REPLACE_AND_RECURE: replace source with an instance of the
                                      quantized type, then try to recursively quantize the child modules of that type
                                    - RECURE_AND_REPLACE: recursively quantize the child modules, then
                                      replace source with an instance of the quantized type
        input_quant_descriptor:     quantization descriptor for inputs (None will take the default one)
        weights_quant_descriptor:   quantization descriptor for weights (None will take the default one)
    """

    class ReplacementAction(Enum):
        REPLACE = "replace"
        REPLACE_AND_RECURE = "replace_and_recure"
        RECURE_AND_REPLACE = "recure_and_replace"
        UNWRAP = "unwrap"
        SKIP = "skip"

    float_source: Union[str, Type]
    quantized_target_class: Optional[Union[Type[QuantMixin], Type[QuantInputMixin], Type[SGQuantMixin]]]
    action: ReplacementAction
    input_quant_descriptor: QuantDescriptor = None  # default is used if None
    weights_quant_descriptor: QuantDescriptor = None  # default is used if None

    def __post_init__(self):
        if self.action in (
            QuantizedMetadata.ReplacementAction.REPLACE,
            QuantizedMetadata.ReplacementAction.REPLACE_AND_RECURE,
            QuantizedMetadata.ReplacementAction.RECURE_AND_REPLACE,
        ):
            assert issubclass(self.quantized_target_class, (SGQuantMixin, QuantMixin, QuantInputMixin))


class QuantizedMapping(nn.Module):
    """
    This class wraps a float module instance, and defines a mapping from this instance to the corresponding quantized
    class, with relevant quant descriptors.

    Example:
        self.my_block = QuantizedMapping(float_module=MyBlock(4, n_classes), quantized_target_class=MyQuantizedBlock)
    """

    def __init__(
        self,
        *,
        float_module: nn.Module,
        quantized_target_class: Union[Type[QuantMixin], Type[QuantInputMixin], Type[SGQuantMixin]],
        action=QuantizedMetadata.ReplacementAction.REPLACE,
        input_quant_descriptor: QuantDescriptor = None,
        weights_quant_descriptor: QuantDescriptor = None,
    ) -> None:
        super().__init__()
        self.float_module = float_module
        self.quantized_target_class = quantized_target_class
        self.action = action
        self.input_quant_descriptor = input_quant_descriptor
        self.weights_quant_descriptor = weights_quant_descriptor
        self.forward = float_module.forward


def _inject_class_methods_to_default_quant_types():
    """
    This is used to add `from_float` capability for the "native" pytorch-quantization (=nvidia-tensorrt) quant classes
    It allows SG to support these modules out of the box
    """
    import pytorch_quantization.quant_modules

    for quant_entry in pytorch_quantization.quant_modules._DEFAULT_QUANT_MAP:
        quant_cls = quant_entry.replace_mod
        quant_cls.from_float = classmethod(_from_float)

import torch

from super_gradients.common.abstractions.abstract_logger import get_logger

logger = get_logger(__name__)

try:
    from pytorch_quantization import nn as quant_nn

    _imported_pytorch_quantization_failure = None
except (ImportError, NameError, ModuleNotFoundError) as import_err:
    logger.warning("Failed to import pytorch_quantization")
    _imported_pytorch_quantization_failure = import_err


def export_quantized_module_to_onnx(model: torch.nn.Module, onnx_filename: str, input_shape: tuple, **kwargs):
    """
    Method for exporting onnx after QAT.

    :param model: torch.nn.Module, model to export
    :param onnx_filename: str, target path for the onnx file,
    :param input_shape: tuple, input shape (usually BCHW)
    """
    if _imported_pytorch_quantization_failure is not None:
        raise _imported_pytorch_quantization_failure

    model.eval()
    if hasattr(model, "prep_model_for_conversion"):
        model.prep_model_for_conversion(**kwargs)

    use_fb_fake_quant_state = quant_nn.TensorQuantizer.use_fb_fake_quant
    quant_nn.TensorQuantizer.use_fb_fake_quant = True

    # Export ONNX for multiple batch sizes
    logger.info("Creating ONNX file: " + onnx_filename)
    dummy_input = torch.randn(input_shape, device=next(model.parameters()).device)
    torch.onnx.export(model, dummy_input, onnx_filename, verbose=False, opset_version=13, enable_onnx_checker=False, do_constant_folding=True)

    # Restore functions of quant_nn back as expected
    quant_nn.TensorQuantizer.use_fb_fake_quant = use_fb_fake_quant_state

from typing import Tuple, Set, Type, Dict, Union, Callable, Optional
from torch import nn

from super_gradients.common.abstractions.abstract_logger import get_logger

logger = get_logger(__name__)
try:
    from pytorch_quantization.nn.modules._utils import QuantMixin, QuantInputMixin
    from pytorch_quantization.tensor_quant import QuantDescriptor
    from pytorch_quantization import nn as quant_nn

    from super_gradients.training.utils.quantization.core import SkipQuantization, SGQuantMixin, QuantizedMapping, QuantizedMetadata

    _imported_pytorch_quantization_failure = None
except (ImportError, NameError, ModuleNotFoundError) as import_err:
    logger.warning("Failed to import pytorch_quantization")
    _imported_pytorch_quantization_failure = import_err


def register_quantized_module(
    float_source: Union[str, Type[nn.Module]],
    action: QuantizedMetadata.ReplacementAction = QuantizedMetadata.ReplacementAction.REPLACE,
    input_quant_descriptor: Optional[QuantDescriptor] = None,
    weights_quant_descriptor: Optional[QuantDescriptor] = None,
) -> Callable:
    """
    Decorator used to register a Quantized module as a quantized version for Float module
    :param action:                      action to perform on the float_source
    :param float_source:                the float module type that is being registered
    :param input_quant_descriptor:      the input quantization descriptor
    :param weights_quant_descriptor:    the weight quantization descriptor
    """

    def decorator(quant_module: Type[SGQuantMixin]) -> Type[SGQuantMixin]:

        if float_source in SelectiveQuantizer.mapping_instructions:
            metadata = SelectiveQuantizer.mapping_instructions[float_source]
            raise ValueError(f"`{float_source}` is already registered with following metadata {metadata}")

        SelectiveQuantizer.mapping_instructions.update(
            {
                float_source: QuantizedMetadata(
                    float_source=float_source,
                    quantized_target_class=quant_module,
                    input_quant_descriptor=input_quant_descriptor,
                    weights_quant_descriptor=weights_quant_descriptor,
                    action=action,
                )
            }
        )
        return quant_module  # this is required since the decorator assigns the result to the `quant_module`

    return decorator


class SelectiveQuantizer:

    """
    :param custom_mappings:                     custom mappings that extend the default mappings with extra behaviour
    :param default_quant_modules_calib_method:  default calibration method (default='percentile')
    :param default_per_channel_quant_modules:   whether quant modules should be per channel (default=False)
    """

    if _imported_pytorch_quantization_failure is not None:
        raise _imported_pytorch_quantization_failure

    mapping_instructions: Dict[Union[str, Type], QuantizedMetadata] = {
        **{
            float_type: QuantizedMetadata(
                float_source=float_type,
                quantized_target_class=quantized_target_class,
                action=QuantizedMetadata.ReplacementAction.REPLACE,
            )
            for (float_type, quantized_target_class) in [
                (nn.Conv1d, quant_nn.QuantConv1d),
                (nn.Conv2d, quant_nn.QuantConv2d),
                (nn.Conv3d, quant_nn.QuantConv3d),
                (nn.ConvTranspose1d, quant_nn.QuantConvTranspose1d),
                (nn.ConvTranspose2d, quant_nn.QuantConvTranspose2d),
                (nn.ConvTranspose3d, quant_nn.QuantConvTranspose3d),
                (nn.Linear, quant_nn.Linear),
                (nn.LSTM, quant_nn.LSTM),
                (nn.LSTMCell, quant_nn.LSTMCell),
                (nn.AvgPool1d, quant_nn.QuantAvgPool1d),
                (nn.AvgPool2d, quant_nn.QuantAvgPool2d),
                (nn.AvgPool3d, quant_nn.QuantAvgPool3d),
                (nn.AdaptiveAvgPool1d, quant_nn.QuantAdaptiveAvgPool1d),
                (nn.AdaptiveAvgPool2d, quant_nn.QuantAdaptiveAvgPool2d),
                (nn.AdaptiveAvgPool3d, quant_nn.QuantAdaptiveAvgPool3d),
            ]
        },
        SkipQuantization: QuantizedMetadata(float_source=SkipQuantization, quantized_target_class=None, action=QuantizedMetadata.ReplacementAction.UNWRAP),
    }  # DEFAULT MAPPING INSTRUCTIONS

    def __init__(
        self, *, custom_mappings: dict = None, default_quant_modules_calib_method: str = "max", default_per_channel_quant_modules: bool = True
    ) -> None:
        super().__init__()
        self.default_quant_modules_calib_method = default_quant_modules_calib_method
        self.default_per_channel_quant_modules = default_per_channel_quant_modules
        self.mapping_instructions = self.mapping_instructions.copy()
        if custom_mappings is not None:
            self.mapping_instructions.update(custom_mappings)  # OVERRIDE DEFAULT WITH CUSTOM. CUSTOM IS PRIORITIZED

    def _get_default_quant_descriptor(self, for_weights=False):
        if self.default_quant_modules_calib_method in ["percentile", "mse", "entropy"]:
            calib_method_type = "histogram"
        else:
            calib_method_type = "max"

        if self.default_per_channel_quant_modules and for_weights:
            return QuantDescriptor(calib_method=calib_method_type, axis=0)
        return QuantDescriptor(calib_method=calib_method_type)

    def register_skip_quantization(self, *, layer_names: Set[str]):
        self.mapping_instructions.update(
            {name: QuantizedMetadata(float_source=name, quantized_target_class=None, action=QuantizedMetadata.ReplacementAction.SKIP) for name in layer_names}
        )

    def register_quantization_mapping(
        self, *, layer_names: Set[str], quantized_target_class: Type[SGQuantMixin], input_quant_descriptor=None, weights_quant_descriptor=None
    ):
        self.mapping_instructions.update(
            {
                name: QuantizedMetadata(
                    float_source=name,
                    quantized_target_class=quantized_target_class,
                    action=QuantizedMetadata.ReplacementAction.REPLACE,
                    input_quant_descriptor=input_quant_descriptor,
                    weights_quant_descriptor=weights_quant_descriptor,
                )
                for name in layer_names
            }
        )

    def _preprocess_skips_and_custom_mappings(self, module: nn.Module, nesting: Tuple[str, ...] = ()):
        """
        This pass is done to extract layer name and mapping instructions, so that we regard to per-layer processing.
        Relevant layer-specific mapping instructions are either `SkipQuantization` or `QuantizedMapping`, which are then
        being added to the mappings
        """
        mapping_instructions = dict()
        for name, child_module in module.named_children():
            nested_name = ".".join(nesting + (name,))
            if isinstance(child_module, SkipQuantization):
                mapping_instructions[nested_name] = QuantizedMetadata(
                    float_source=nested_name, quantized_target_class=None, action=QuantizedMetadata.ReplacementAction.UNWRAP
                )

            if isinstance(child_module, QuantizedMapping):
                mapping_instructions[nested_name] = QuantizedMetadata(
                    float_source=nested_name,
                    quantized_target_class=child_module.quantized_target_class,
                    input_quant_descriptor=child_module.input_quant_descriptor,
                    weights_quant_descriptor=child_module.weights_quant_descriptor,
                    action=child_module.action,
                )

            if isinstance(child_module, nn.Module):  # recursive call
                mapping_instructions.update(self._preprocess_skips_and_custom_mappings(child_module, nesting + (name,)))

        return mapping_instructions

    def _instantiate_quantized_from_float(self, float_module, metadata, preserve_state_dict):
        base_classes = (QuantMixin, QuantInputMixin, SGQuantMixin)
        if not issubclass(metadata.quantized_target_class, base_classes):
            raise AssertionError(
                f"Quantization suite for {type(float_module).__name__} is invalid. "
                f"{metadata.quantized_target_class.__name__} must inherit one of "
                f"{', '.join(map(lambda _: _.__name__, base_classes))}"
            )

        # USE PROVIDED QUANT DESCRIPTORS, OR DEFAULT IF NONE PROVIDED
        quant_descriptors = dict()
        if issubclass(metadata.quantized_target_class, (SGQuantMixin, QuantMixin, QuantInputMixin)):
            quant_descriptors = {"quant_desc_input": metadata.input_quant_descriptor or self._get_default_quant_descriptor(for_weights=False)}
        if issubclass(metadata.quantized_target_class, (SGQuantMixin, QuantMixin)):
            quant_descriptors.update({"quant_desc_weight": metadata.weights_quant_descriptor or self._get_default_quant_descriptor(for_weights=True)})

        if not hasattr(metadata.quantized_target_class, "from_float"):
            assert isinstance(metadata.quantized_target_class, SGQuantMixin), (
                f"{metadata.quantized_target_class.__name__} must inherit from " f"{SGQuantMixin.__name__}, so that it would include `from_float` class method"
            )

        q_instance = metadata.quantized_target_class.from_float(float_module, **quant_descriptors)

        # MOVE TENSORS TO ORIGINAL DEVICE
        if len(list(float_module.parameters(recurse=False))) > 0:
            q_instance = q_instance.to(next(float_module.parameters(recurse=False)).device)
        elif len(list(float_module.buffers(recurse=False))):
            q_instance = q_instance.to(next(float_module.buffers(recurse=False)).device)

        # COPY STATE DICT IF NEEDED
        if preserve_state_dict:
            q_instance.load_state_dict(float_module.state_dict(), strict=True)

        return q_instance

    def _maybe_quantize_one_layer(
        self,
        module: nn.Module,
        child_name: str,
        nesting: Tuple[str, ...],
        child_module: nn.Module,
        mapping_instructions: Dict[Union[str, Type], QuantizedMetadata],
        preserve_state_dict: bool,
    ) -> bool:
        """
        Does the heavy lifting of (maybe) quantizing a layer: creates a quantized instance based on a float instance,
        and replaces it in the "parent" module

        :param module:                  the module we'd like to quantize a specific layer in
        :param child_name:              the attribute name of the layer in the module
        :param nesting:                 the current nesting we're in. Needed to find the appropriate key in the mappings
        :param child_module:            the instance of the float module we'd like to quantize
        :param mapping_instructions:    mapping instructions: how to quantize
        :param preserve_state_dict:     whether to copy the state dict from the float instance to the quantized instance

        :return: a boolean indicates if we found a match and should not continue recursively
        """
        # if we don't have any instruction for the specific layer or the specific type - we continue
        # NOTE! IT IS IMPORTANT TO FIRST PROCESS THE NAME AND ONLY THEN THE TYPE
        if _imported_pytorch_quantization_failure is not None:
            raise _imported_pytorch_quantization_failure
        for candidate_key in (".".join(nesting + (child_name,)), type(child_module)):
            if candidate_key not in mapping_instructions:
                continue
            metadata: QuantizedMetadata = mapping_instructions[candidate_key]
            if metadata.action == QuantizedMetadata.ReplacementAction.SKIP:
                return True
            elif metadata.action == QuantizedMetadata.ReplacementAction.UNWRAP:
                assert isinstance(child_module, SkipQuantization)
                setattr(module, child_name, child_module.float_module)
                return True
            elif metadata.action in (
                QuantizedMetadata.ReplacementAction.REPLACE,
                QuantizedMetadata.ReplacementAction.REPLACE_AND_RECURE,
                QuantizedMetadata.ReplacementAction.RECURE_AND_REPLACE,
            ):
                if isinstance(child_module, QuantizedMapping):  # UNWRAP MAPPING
                    child_module = child_module.float_module
                q_instance: nn.Module = self._instantiate_quantized_from_float(
                    float_module=child_module, metadata=metadata, preserve_state_dict=preserve_state_dict
                )

                # ACTUAL REPLACEMENT
                def replace():
                    setattr(module, child_name, q_instance)

                def recurse_quantize():
                    self._quantize_module_aux(
                        module=getattr(module, child_name),
                        mapping_instructions=mapping_instructions,
                        nesting=nesting + (child_name,),
                        preserve_state_dict=preserve_state_dict,
                    )

                if metadata.action == QuantizedMetadata.ReplacementAction.REPLACE:
                    replace()
                elif metadata.action == QuantizedMetadata.ReplacementAction.REPLACE_AND_RECURE:
                    replace()
                    recurse_quantize()
                elif metadata.action == QuantizedMetadata.ReplacementAction.RECURE_AND_REPLACE:
                    recurse_quantize()
                    replace()
                return True
            else:
                raise NotImplementedError
        return False

    def quantize_module(self, module: nn.Module, *, preserve_state_dict=True):
        per_layer_mappings = self._preprocess_skips_and_custom_mappings(module)
        mapping_instructions = {
            **per_layer_mappings,
            **self.mapping_instructions,
        }  # we first regard the per layer mappings, and then override with the custom mappings in case there is overlap

        self._quantize_module_aux(mapping_instructions=mapping_instructions, module=module, nesting=(), preserve_state_dict=preserve_state_dict)

    def _quantize_module_aux(self, mapping_instructions, module, nesting, preserve_state_dict):
        for name, child_module in module.named_children():
            found = self._maybe_quantize_one_layer(module, name, nesting, child_module, mapping_instructions, preserve_state_dict)

            # RECURSIVE CALL, to support module_list, sequential, custom (nested) modules
            if not found and isinstance(child_module, nn.Module):
                self._quantize_module_aux(mapping_instructions, child_module, nesting + (name,), preserve_state_dict)

"""
Quantization utilities

Methods are based on:
 https://github.com/NVIDIA/TensorRT/blob/51a4297753d3e12d0eed864be52400f429a6a94c/tools/pytorch-quantization/examples/torchvision/classification_flow.py#L385

(Licensed under the Apache License, Version 2.0)
"""
from torch.utils.data import DataLoader
from tqdm import tqdm
import torch
from super_gradients.common.abstractions.abstract_logger import get_logger
from super_gradients.training import models
from super_gradients.training.utils.callbacks import Phase, PhaseCallback, PhaseContext
import os
from enum import Enum
from super_gradients.training.utils.checkpoint_utils import load_checkpoint_to_model
from super_gradients.training.utils import get_param
from super_gradients.training.utils.distributed_training_utils import get_local_rank, get_world_size
from torch.distributed import all_gather

logger = get_logger(__name__)

try:
    from pytorch_quantization import nn as quant_nn, quant_modules
    from pytorch_quantization import calib
    from pytorch_quantization.tensor_quant import QuantDescriptor

    _imported_pytorch_quantization_failure = None
except (ImportError, NameError, ModuleNotFoundError) as import_err:
    logger.debug("Failed to import pytorch_quantization")
    _imported_pytorch_quantization_failure = import_err


class QuantizationLevel(str, Enum):
    FP32 = "FP32"
    FP16 = "FP16"
    INT8 = "INT8"
    HYBRID = "Hybrid"

    @staticmethod
    def from_string(quantization_level: str) -> Enum:
        quantization_level = quantization_level.lower()
        if quantization_level == "fp32":
            return QuantizationLevel.FP32
        elif quantization_level == "fp16":
            return QuantizationLevel.FP16
        elif quantization_level == "int8":
            return QuantizationLevel.INT8
        elif quantization_level == "hybrid":
            return QuantizationLevel.HYBRID
        else:
            raise NotImplementedError(f'Quantization Level: "{quantization_level}" is not supported')


def export_qat_onnx(model: torch.nn.Module, onnx_filename: str, input_shape: tuple, per_channel_quantization: bool = False):
    """
    Method for exporting onnx after QAT.

    :param model: torch.nn.Module, model to export
    :param onnx_filename: str, target path for the onnx file,
    :param input_shape: tuple, input shape (usually BCHW)
    """
    if _imported_pytorch_quantization_failure is not None:
        raise _imported_pytorch_quantization_failure
    else:
        model.eval()
        if hasattr(model, "prep_model_for_conversion"):
            model.prep_model_for_conversion()
        quant_nn.TensorQuantizer.use_fb_fake_quant = True
        # Export ONNX for multiple batch sizes
        logger.info("Creating ONNX file: " + onnx_filename)
        dummy_input = torch.randn(input_shape, device="cuda")
        opset_version = 13 if per_channel_quantization else 12
        torch.onnx.export(model, dummy_input, onnx_filename, verbose=False, opset_version=opset_version, enable_onnx_checker=False, do_constant_folding=True)


def calibrate_model(
    model: torch.nn.Module, calib_data_loader: torch.utils.data.DataLoader, method: str = "percentile", num_calib_batches: int = 2, percentile: float = 99.99
):
    """
    Calibrates torch model with quantized modules.

    :param model:               torch.nn.Module, model to perfrom the calibration on.
    :param calib_data_loader:   torch.utils.data.DataLoader, data loader of the calibration dataset.
    :param method:              str, One of [percentile, mse, entropy, max]. Statistics method for amax computation of the quantized modules
                                (Default=percentile).
    :param num_calib_batches:   int, number of batches to collect the statistics from.
    :param percentile:          float, percentile value to use when Trainer,quant_modules_calib_method='percentile'. Discarded when other methods are used
                                (Default=99.99).

    """
    if _imported_pytorch_quantization_failure is not None:
        raise _imported_pytorch_quantization_failure
    elif method in ["percentile", "mse", "entropy", "max"]:
        with torch.no_grad():
            _collect_stats(model, calib_data_loader, num_batches=num_calib_batches)

            # FOR PERCENTILE WE MUST PASS PERCENTILE VALUE THROUGH KWARGS,
            # SO IT WOULD BE PASSED TO module.load_calib_amax(**kwargs), AND IN OTHER METHODS WE MUST NOT PASS IT.
            if method == "precentile":
                _compute_amax(model, method="percentile", percentile=percentile)
            else:
                _compute_amax(model, method=method)
    else:
        raise ValueError("Unsupported quantization calibration method, expected one of: percentile, mse, entropy, max, got " + str(method) + ".")


def _collect_stats(model, data_loader, num_batches):
    """Feed data to the network and collect statistics"""
    if _imported_pytorch_quantization_failure is not None:
        raise _imported_pytorch_quantization_failure
    else:
        local_rank = get_local_rank()
        world_size = get_world_size()

        # Enable calibrators
        _enable_calibrators(model)

        # Feed data to the network for collecting stats
        for i, (image, _) in tqdm(enumerate(data_loader), total=num_batches, disable=local_rank > 0):
            if world_size > 1:
                all_batches = [torch.zeros_like(image, device="cuda") for _ in range(world_size)]
                all_gather(all_batches, image.cuda())
            else:
                all_batches = [image]

            for local_image in all_batches:
                model(local_image.cuda())
            if i >= num_batches:
                break

        # Disable calibrators
        _disable_calibrators(model)


def _disable_calibrators(model):
    for name, module in model.named_modules():
        if isinstance(module, quant_nn.TensorQuantizer):
            if module._calibrator is not None:
                module.enable_quant()
                module.disable_calib()
            else:
                module.enable()


def _enable_calibrators(model):
    for name, module in model.named_modules():
        if isinstance(module, quant_nn.TensorQuantizer):
            if module._calibrator is not None:
                module.disable_quant()
                module.enable_calib()
            else:
                module.disable()


def _compute_amax(model, **kwargs):
    if _imported_pytorch_quantization_failure is not None:
        raise _imported_pytorch_quantization_failure
    else:
        # Load calib result
        for name, module in model.named_modules():
            if isinstance(module, quant_nn.TensorQuantizer):
                if module._calibrator is not None:
                    if isinstance(module._calibrator, calib.MaxCalibrator):
                        module.load_calib_amax()
                    else:
                        module.load_calib_amax(**kwargs)
        model.cuda()


def _deactivate_quant_modules_wrapping():
    """
    Deactivates quant modules wrapping, so that further modules won't use Q/DQ layers.
    """
    if _imported_pytorch_quantization_failure is not None:
        raise _imported_pytorch_quantization_failure
    else:
        quant_modules.deactivate()


def _activate_quant_modules_wrapping():
    """
    Activates quant modules wrapping, so that further modules use Q/DQ layers.
    """
    if _imported_pytorch_quantization_failure is not None:
        raise _imported_pytorch_quantization_failure
    else:
        quant_modules.initialize()


class QATCallback(PhaseCallback):
    """
    A callback for transitioning training into QAT.

    Rebuilds the model with QAT layers then either:
        1. loads the best checkpoint then performs calibration.
        2. loads an external calibrated model (makes sense when start_epoch=0).

    Additionally, resets Trainer's best_metric and sets ckpt_best_name to 'qat_ckpt_best.pth' so best QAT checkpoints
     will be saved separately.

    If performing calibration- the calibrated model is evaluated, and the metric_to_watch is logged under
     calibrated_model_{metric_to_watch}. The calibrated checkpoint is saved under ckpt_calibrated_{calibration_method}.pth


    Attributes:
        start_epoch: int, first epoch to start QAT.

        quant_modules_calib_method: str, One of [percentile, mse, entropy, max]. Statistics method for amax
         computation of the quantized modules (default=percentile).

        per_channel_quant_modules: bool, whether quant modules should be per channel (default=False).

        calibrate: bool, whether to perfrom calibration (default=False).

        calibrated_model_path: str, path to a calibrated checkpoint (default=None).

        calib_data_loader: torch.utils.data.DataLoader, data loader of the calibration dataset. When None,
         context.train_loader will be used (default=None).

        num_calib_batches: int, number of batches to collect the statistics from.

        percentile: float, percentile value to use when Trainer,quant_modules_calib_method='percentile'.
         Discarded when other methods are used (Default=99.99).



    """

    def __init__(
        self,
        start_epoch: int,
        quant_modules_calib_method: str = "percentile",
        per_channel_quant_modules: bool = False,
        calibrate: bool = True,
        calibrated_model_path: str = None,
        calib_data_loader: DataLoader = None,
        num_calib_batches: int = 2,
        percentile: float = 99.99,
    ):
        super(QATCallback, self).__init__(Phase.TRAIN_EPOCH_START)
        self._validate_args(start_epoch, quant_modules_calib_method, calibrate, calibrated_model_path)
        self.start_epoch = start_epoch
        self.quant_modules_calib_method = quant_modules_calib_method
        self.per_channel_quant_modules = per_channel_quant_modules
        self.calibrate = calibrate
        self.calibrated_model_path = calibrated_model_path
        self.calib_data_loader = calib_data_loader
        self.num_calib_batches = num_calib_batches
        self.percentile = percentile

    def _validate_args(self, start_epoch: int, quant_modules_calib_method: str, calibrate, calibrated_model_path):
        if _imported_pytorch_quantization_failure:
            raise _imported_pytorch_quantization_failure
        if start_epoch < 0:
            raise ValueError("start_epoch must be positive.")
        if quant_modules_calib_method not in ["percentile", "mse", "entropy", "max"]:
            raise ValueError(
                "Unsupported quantization calibration method, expected one of: percentile, mse, entropy, max, got " + str(self.quant_modules_calib_method) + "."
            )
        if not calibrate and calibrated_model_path is None:
            logger.warning("calibrate=False and no calibrated_model_path is given. QAT will be on an uncalibrated model.")

    def __call__(self, context: PhaseContext):
        if context.epoch == self.start_epoch:
            # REMOVE REFERENCES TO NETWORK AND CLEAN GPU MEMORY BEFORE BUILDING THE NEW NET
            context.context_methods.set_net(None)
            context.net = None
            torch.cuda.empty_cache()

            # BUILD THE SAME MODEL BUT WITH FAKE QUANTIZED MODULES, AND LOAD BEST CHECKPOINT TO IT
            self._initialize_quant_modules()

            if self.calibrated_model_path is not None:
                checkpoint_path = self.calibrated_model_path
            elif self.start_epoch > 0:
                checkpoint_path = os.path.join(context.ckpt_dir, "ckpt_best.pth")

            qat_net = models.get(context.architecture, arch_params=context.arch_params.to_dict(), checkpoint_path=checkpoint_path)

            _deactivate_quant_modules_wrapping()

            # UPDATE CONTEXT'S NET REFERENCE
            context.net = context.context_methods.get_net()

            if self.calibrate:
                self._calibrate_model(context)

            # RESET THE BEST METRIC VALUE SO WE SAVE CHECKPOINTS AFTER THE EXPECTED QAT ACCURACY DEGRADATION
            context.context_methods._reset_best_metric()

            # SET NEW FILENAME FOR THE BEST CHECKPOINT SO WE DON'T OVERRIDE THE PREVIOUS ONES
            context.context_methods.set_ckpt_best_name("qat_ckpt_best.pth")

            # FINALLY, SET THE QAT NET TO CONTINUE TRAINING
            context.context_methods.set_net(qat_net)

    def _calibrate_model(self, context: PhaseContext):
        """
        Performs model calibration (collecting stats and setting amax for the fake quantized moduls)

        :param context: PhaseContext, current phase context.
        """
        self.calib_data_loader = self.calib_data_loader or context.train_loader
        calibrate_model(
            model=context.net,
            calib_data_loader=self.calib_data_loader,
            method=self.quant_modules_calib_method,
            num_calib_batches=self.num_calib_batches,
            percentile=self.percentile,
        )
        method_desc = (
            self.quant_modules_calib_method + "_" + str(self.percentile) if self.quant_modules_calib_method == "percentile" else self.quant_modules_calib_method
        )

        if not context.ddp_silent_mode:
            logger.info("Performing additional validation on calibrated model...")

        calibrated_valid_results = context.context_methods.validate_epoch(epoch=self.start_epoch, silent_mode=True)
        calibrated_acc = calibrated_valid_results[context.metric_idx_in_results_tuple]

        if not context.ddp_silent_mode:
            logger.info("Calibrate model " + context.metric_to_watch + ": " + str(calibrated_acc))
            context.sg_logger.add_checkpoint(tag="ckpt_calibrated_" + method_desc + ".pth", state_dict={"net": context.net.state_dict(), "acc": calibrated_acc})
            context.sg_logger.add_scalar("Calibrated_Model_" + context.metric_to_watch, calibrated_acc, global_step=self.start_epoch)

    def _initialize_quant_modules(self):
        """
        Initialize quant modules wrapping.
        """

        if _imported_pytorch_quantization_failure is not None:
            raise _imported_pytorch_quantization_failure
        else:
            if self.quant_modules_calib_method in ["percentile", "mse", "entropy"]:
                calib_method_type = "histogram"
            else:
                calib_method_type = "max"

            if self.per_channel_quant_modules:
                quant_desc_input = QuantDescriptor(calib_method=calib_method_type)
                quant_nn.QuantConv2d.set_default_quant_desc_input(quant_desc_input)
                quant_nn.QuantLinear.set_default_quant_desc_input(quant_desc_input)
            else:
                quant_desc_input = QuantDescriptor(calib_method=calib_method_type, axis=None)
                quant_nn.QuantConv2d.set_default_quant_desc_input(quant_desc_input)
                quant_nn.QuantConvTranspose2d.set_default_quant_desc_input(quant_desc_input)
                quant_nn.QuantLinear.set_default_quant_desc_input(quant_desc_input)

                quant_desc_weight = QuantDescriptor(calib_method=calib_method_type, axis=None)
                quant_nn.QuantConv2d.set_default_quant_desc_weight(quant_desc_weight)
                quant_nn.QuantConvTranspose2d.set_default_quant_desc_weight(quant_desc_weight)
                quant_nn.QuantLinear.set_default_quant_desc_weight(quant_desc_weight)

            _activate_quant_modules_wrapping()


class PostQATConversionCallback(PhaseCallback):
    """
    Post QAT training callback that saves the best checkpoint (i.e ckpt_best.pth) in onnx format.
    Should be used with QATCallback.

    Attributes:
        dummy_input_size: (tuple) dummy input size for the ONNX conversion.
    """

    def __init__(self, dummy_input_size):
        super().__init__(phase=Phase.POST_TRAINING)
        self.dummy_input_size = dummy_input_size

    def __call__(self, context: PhaseContext):
        if not context.ddp_silent_mode:
            best_ckpt_path = os.path.join(context.ckpt_dir, "qat_ckpt_best.pth")
            onnx_path = os.path.join(context.ckpt_dir, "qat_ckpt_best.onnx")

            load_checkpoint_to_model(
                ckpt_local_path=best_ckpt_path,
                net=context.net,
                load_weights_only=True,
                load_ema_as_net=context.training_params.ema,
                strict=True,
                load_backbone=False,
            )
            per_channel_quant_modules = get_param(context.training_params.qat_params, "per_channel_quant_modules")
            export_qat_onnx(context.net.module, onnx_path, self.dummy_input_size, per_channel_quant_modules)

            context.sg_logger.add_file("qat_ckpt_best.onnx")

@@ -22,6 +22,7 @@ from tests.unit_tests import (
 )
 from tests.end_to_end_tests import TestTrainer
 from tests.unit_tests.phase_delegates_test import ContextMethodsTest
+from tests.unit_tests.quantization_utility_tests import QuantizationUtilityTest
 from tests.unit_tests.random_erase_test import RandomEraseTest
 from tests.unit_tests.strictload_enum_test import StrictLoadEnumTest
 from tests.unit_tests.train_with_intialized_param_args_test import TrainWithInitializedObjectsTest
@@ -90,6 +91,7 @@ class CoreUnitTestSuiteRunner:
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(IoULossTest))
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(TestDetectionDatasetSubsampling))
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(TestDetectionDatasetSubclassing))
+        self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(QuantizationUtilityTest))
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(TestDetectionDatasetCaching))
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(MultiScaleTest))
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(TrainingParamsTest))

import unittest

from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader
from super_gradients.training import Trainer, MultiGPUMode, models
from super_gradients.training.metrics.classification_metrics import Accuracy
import os
from super_gradients.training.utils.quantization_utils import PostQATConversionCallback


class QATIntegrationTest(unittest.TestCase):
    def _get_trainer(self, experiment_name):
        trainer = Trainer(experiment_name,

                          multi_gpu=MultiGPUMode.OFF)
        model = models.get("resnet18", pretrained_weights="imagenet")
        return trainer, model

    def _get_train_params(self, qat_params):
        train_params = {"max_epochs": 2,
                        "lr_mode": "step",
                        "optimizer": "SGD",
                        "lr_updates": [],
                        "lr_decay_factor": 0.1,
                        "initial_lr": 0.001, "loss": "cross_entropy",
                        "train_metrics_list": [Accuracy()],
                        "valid_metrics_list": [Accuracy()],

                        "metric_to_watch": "Accuracy",
                        "greater_metric_to_watch_is_better": True,
                        "average_best_models": False,
                        "enable_qat": True,
                        "qat_params": qat_params,
                        "phase_callbacks": [PostQATConversionCallback(dummy_input_size=(1, 3, 224, 224))]
                        }
        return train_params

    def test_qat_from_start(self):
        model, net = self._get_trainer("test_qat_from_start")
        train_params = self._get_train_params(qat_params={
            "start_epoch": 0,
            "quant_modules_calib_method": "percentile",
            "calibrate": True,
            "num_calib_batches": 2,
            "percentile": 99.99
        })

        model.train(model=net, training_params=train_params, train_loader=classification_test_dataloader(),
                    valid_loader=classification_test_dataloader())

    def test_qat_transition(self):
        model, net = self._get_trainer("test_qat_transition")
        train_params = self._get_train_params(qat_params={
            "start_epoch": 1,
            "quant_modules_calib_method": "percentile",
            "calibrate": True,
            "num_calib_batches": 2,
            "percentile": 99.99
        })

        model.train(model=net, training_params=train_params, train_loader=classification_test_dataloader(),
                    valid_loader=classification_test_dataloader())

    def test_qat_from_calibrated_ckpt(self):
        model, net = self._get_trainer("generate_calibrated_model")
        train_params = self._get_train_params(qat_params={
            "start_epoch": 0,
            "quant_modules_calib_method": "percentile",
            "calibrate": True,
            "num_calib_batches": 2,
            "percentile": 99.99
        })

        model.train(model=net, training_params=train_params, train_loader=classification_test_dataloader(),
                    valid_loader=classification_test_dataloader())

        calibrated_model_path = os.path.join(model.checkpoints_dir_path, "ckpt_calibrated_percentile_99.99.pth")

        model, net = self._get_trainer("test_qat_from_calibrated_ckpt")
        train_params = self._get_train_params(qat_params={
            "start_epoch": 0,
            "quant_modules_calib_method": "percentile",
            "calibrate": False,
            "calibrated_model_path": calibrated_model_path,
            "num_calib_batches": 2,
            "percentile": 99.99
        })

        model.train(model=net, training_params=train_params, train_loader=classification_test_dataloader(),
                    valid_loader=classification_test_dataloader())


if __name__ == '__main__':
    unittest.main()

import unittest
import torch
import torchvision
from torch import nn

try:
    import super_gradients
    from pytorch_quantization import nn as quant_nn
    from pytorch_quantization import quant_modules
    from super_gradients.training.utils.quantization.selective_quantization_utils import SelectiveQuantizer, register_quantized_module
    from pytorch_quantization.calib import MaxCalibrator, HistogramCalibrator
    from super_gradients.training.utils.quantization.core import SkipQuantization, SGQuantMixin, QuantizedMapping, QuantizedMetadata
    from pytorch_quantization.nn import QuantConv2d
    from pytorch_quantization.tensor_quant import QuantDescriptor

    _imported_pytorch_quantization_failure = False

except (ImportError, NameError, ModuleNotFoundError):
    _imported_pytorch_quantization_failure = True


@unittest.skipIf(_imported_pytorch_quantization_failure, "Failed to import `pytorch_quantization`")
class QuantizationUtilityTest(unittest.TestCase):
    def test_vanilla_replacement(self):
        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)

            def forward(self, x):
                return self.conv1(x)

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer()
        q_util.quantize_module(module)

        x = torch.rand(1, 3, 32, 32)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, 8, 32, 32))

        self.assertTrue(isinstance(module.conv1, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))

    def test_module_list_replacement(self):
        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.convs = nn.ModuleList([nn.Conv2d(3, 8, kernel_size=3, padding=1) for _ in range(3)])

            def forward(self, x):
                return torch.cat([conv(x) for conv in self.convs], dim=1)

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer()
        q_util.quantize_module(module)

        x = torch.rand(1, 3, 32, 32)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, 3 * 8, 32, 32))

        for conv in module.convs:
            self.assertTrue(isinstance(conv, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))

    def test_sequential_list_replacement(self):
        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.convs = nn.Sequential(
                    nn.Conv2d(3, 8, kernel_size=3, padding=1),
                    nn.Conv2d(8, 16, kernel_size=3, padding=1),
                )

            def forward(self, x):
                return self.convs(x)

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer()
        q_util.quantize_module(module)

        x = torch.rand(1, 3, 32, 32)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, 16, 32, 32))

        for conv in module.convs:
            self.assertTrue(isinstance(conv, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))

    def test_nested_module_replacement(self):
        # ARRANGE
        class MyBlock(nn.Module):
            def __init__(self, in_feats, out_feats) -> None:
                super().__init__()
                self.flatten = nn.Flatten()
                self.linear = nn.Linear(in_feats, out_feats)

            def forward(self, x):
                return self.linear(self.flatten(x))

        class MyModel(nn.Module):
            def __init__(self, res, n_classes) -> None:
                super().__init__()
                self.conv = nn.Conv2d(3, 4, kernel_size=3, padding=1)
                self.my_block = MyBlock(4 * (res**2), n_classes)

            def forward(self, x):
                y = self.conv(x)
                return self.my_block(y)

        res = 32
        n_clss = 10
        module = MyModel(res, n_clss)

        # TEST
        q_util = SelectiveQuantizer()
        q_util.quantize_module(module)

        x = torch.rand(1, 3, res, res)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, n_clss))

        self.assertTrue(isinstance(module.conv, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(isinstance(module.my_block.linear, SelectiveQuantizer.mapping_instructions[nn.Linear].quantized_target_class))

    def test_static_selective_skip_quantization(self):
        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)
                self.conv2 = SkipQuantization(nn.Conv2d(8, 16, kernel_size=3, padding=1))

            def forward(self, x):
                return self.conv2(self.conv1(x))

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer()
        q_util.quantize_module(module)

        x = torch.rand(1, 3, 32, 32)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, 16, 32, 32))

        self.assertTrue(isinstance(module.conv1, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(isinstance(module.conv2, nn.Conv2d))

    def test_dynamic_skip_quantization(self):
        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)
                self.conv2 = nn.Conv2d(8, 16, kernel_size=3, padding=1)

            def forward(self, x):
                return self.conv2(self.conv1(x))

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer()
        q_util.register_skip_quantization(layer_names={"conv2"})
        q_util.quantize_module(module)
        x = torch.rand(1, 3, 32, 32)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, 16, 32, 32))

        self.assertTrue(isinstance(module.conv1, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(isinstance(module.conv2, nn.Conv2d))

    def test_custom_quantized_mapping_wrapper_explicit_from_float(self):
        # ARRANGE
        class MyBlock(nn.Module):
            def __init__(self, in_feats, out_feats) -> None:
                super().__init__()
                self.flatten = nn.Flatten()
                self.linear = nn.Linear(in_feats, out_feats)

            def forward(self, x):
                return self.linear(self.flatten(x))

        class MyQuantizedBlock(SGQuantMixin):
            # NOTE: **kwargs are necessary because quant descriptors are passed there!
            @classmethod
            def from_float(cls, float_instance: MyBlock, **kwargs):
                return cls(in_feats=float_instance.linear.in_features, out_feats=float_instance.linear.out_features)

            def __init__(self, in_feats, out_feats) -> None:
                super().__init__()
                self.flatten = nn.Flatten()
                self.linear = quant_nn.QuantLinear(in_feats, out_feats)

            def forward(self, x):
                return self.linear(self.flatten(x))

        class MyModel(nn.Module):
            def __init__(self, res, n_classes) -> None:
                super().__init__()
                self.conv = nn.Conv2d(3, 4, kernel_size=3, padding=1)
                self.my_block = QuantizedMapping(float_module=MyBlock(4 * (res**2), n_classes), quantized_target_class=MyQuantizedBlock)

            def forward(self, x):
                y = self.conv(x)
                return self.my_block(y)

        res = 32
        n_clss = 10
        module = MyModel(res, n_clss)

        # TEST
        q_util = SelectiveQuantizer()
        q_util.quantize_module(module)

        x = torch.rand(1, 3, res, res)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, n_clss))

        self.assertTrue(isinstance(module.conv, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(isinstance(module.my_block, MyQuantizedBlock))

    def test_custom_quantized_mapping_wrapper_implicit_from_float(self):
        # ARRANGE
        class MyBlock(nn.Module):
            def __init__(self, in_feats, out_feats) -> None:
                super().__init__()
                self.in_feats = in_feats
                self.out_feats = out_feats
                self.flatten = nn.Flatten()
                self.linear = nn.Linear(in_feats, out_feats)

            def forward(self, x):
                return self.linear(self.flatten(x))

        class MyQuantizedBlock(SGQuantMixin):
            def __init__(self, in_feats, out_feats) -> None:
                super().__init__()
                self.flatten = nn.Flatten()
                self.linear = quant_nn.QuantLinear(in_feats, out_feats)

            def forward(self, x):
                return self.linear(self.flatten(x))

        class MyModel(nn.Module):
            def __init__(self, res, n_classes) -> None:
                super().__init__()
                self.conv = nn.Conv2d(3, 4, kernel_size=3, padding=1)
                self.my_block = QuantizedMapping(float_module=MyBlock(4 * (res**2), n_classes), quantized_target_class=MyQuantizedBlock)

            def forward(self, x):
                y = self.conv(x)
                return self.my_block(y)

        res = 32
        n_clss = 10
        module = MyModel(res, n_clss)

        # TEST
        q_util = SelectiveQuantizer()
        q_util.quantize_module(module)

        x = torch.rand(1, 3, res, res)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, n_clss))

        self.assertTrue(isinstance(module.conv, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(isinstance(module.my_block, MyQuantizedBlock))

    def test_custom_quantized_mapping_register_with_decorator(self):
        # ARRANGE
        class MyBlock(nn.Module):
            def __init__(self, in_feats, out_feats) -> None:
                super().__init__()
                self.in_feats = in_feats
                self.out_feats = out_feats
                self.flatten = nn.Flatten()
                self.linear = nn.Linear(in_feats, out_feats)

            def forward(self, x):
                return self.linear(self.flatten(x))

        @register_quantized_module(float_source=MyBlock)
        class MyQuantizedBlock(SGQuantMixin):
            def __init__(self, in_feats, out_feats) -> None:
                super().__init__()
                self.flatten = nn.Flatten()
                self.linear = quant_nn.QuantLinear(in_feats, out_feats)

            def forward(self, x):
                return self.linear(self.flatten(x))

        class MyModel(nn.Module):
            def __init__(self, res, n_classes) -> None:
                super().__init__()
                self.conv = nn.Conv2d(3, 4, kernel_size=3, padding=1)
                self.my_block = MyBlock(4 * (res**2), n_classes)

            def forward(self, x):
                y = self.conv(x)
                return self.my_block(y)

        res = 32
        n_clss = 10
        module = MyModel(res, n_clss)

        # TEST
        q_util = SelectiveQuantizer()
        q_util.quantize_module(module)

        x = torch.rand(1, 3, res, res)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, n_clss))
        self.assertTrue(MyQuantizedBlock is not None)
        self.assertTrue(isinstance(module.conv, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(isinstance(module.my_block, MyQuantizedBlock))

    def test_dynamic_quantized_mapping(self):
        # ARRANGE
        class MyBlock(nn.Module):
            def __init__(self, in_feats, out_feats) -> None:
                super().__init__()
                self.in_feats = in_feats
                self.out_feats = out_feats
                self.flatten = nn.Flatten()
                self.linear = nn.Linear(in_feats, out_feats)

            def forward(self, x):
                return self.linear(self.flatten(x))

        class MyQuantizedBlock(SGQuantMixin):
            def __init__(self, in_feats, out_feats) -> None:
                super().__init__()
                self.flatten = nn.Flatten()
                self.linear = quant_nn.QuantLinear(in_feats, out_feats)

            def forward(self, x):
                return self.linear(self.flatten(x))

        class MyModel(nn.Module):
            def __init__(self, res, n_classes) -> None:
                super().__init__()
                self.conv = nn.Conv2d(3, 4, kernel_size=3, padding=1)
                self.my_block = MyBlock(4 * (res**2), n_classes)

            def forward(self, x):
                y = self.conv(x)
                return self.my_block(y)

        res = 32
        n_clss = 10
        module = MyModel(res, n_clss)

        # TEST
        q_util = SelectiveQuantizer()
        q_util.register_quantization_mapping(layer_names={"my_block"}, quantized_target_class=MyQuantizedBlock)
        q_util.quantize_module(module)

        x = torch.rand(1, 3, res, res)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, n_clss))

        self.assertTrue(isinstance(module.conv, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(isinstance(module.my_block, MyQuantizedBlock))

    def test_non_default_quant_descriptors_are_piped(self):
        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)

            def forward(self, x):
                return self.conv1(x)

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer(default_quant_modules_calib_method="max")
        q_util.quantize_module(module)

        x = torch.rand(1, 3, 32, 32)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, 8, 32, 32))
        self.assertTrue(isinstance(module.conv1, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(type(module.conv1._input_quantizer._calibrator) == MaxCalibrator)
        self.assertTrue(type(module.conv1._weight_quantizer._calibrator) == MaxCalibrator)

    def test_different_quant_descriptors_are_piped(self):
        # ARRANGE
        class MyModel(nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)
                self.conv2 = nn.Conv2d(8, 8, kernel_size=3, padding=1)

            def forward(self, x):
                return self.conv2(self.conv1(x))

        module = MyModel()

        # TEST
        q_util = SelectiveQuantizer()
        q_util.register_quantization_mapping(
            layer_names={"conv1"},
            quantized_target_class=QuantConv2d,
            input_quant_descriptor=QuantDescriptor(calib_method="max"),
            weights_quant_descriptor=QuantDescriptor(calib_method="histogram"),
        )
        q_util.register_quantization_mapping(
            layer_names={"conv2"},
            quantized_target_class=QuantConv2d,
            input_quant_descriptor=QuantDescriptor(calib_method="histogram"),
            weights_quant_descriptor=QuantDescriptor(calib_method="max"),
        )
        q_util.quantize_module(module)

        x = torch.rand(1, 3, 32, 32)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, 8, 32, 32))
        self.assertTrue(isinstance(module.conv1, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(type(module.conv1._input_quantizer._calibrator) == MaxCalibrator)
        self.assertTrue(type(module.conv1._weight_quantizer._calibrator) == HistogramCalibrator)
        self.assertTrue(type(module.conv2._input_quantizer._calibrator) == HistogramCalibrator)
        self.assertTrue(type(module.conv2._weight_quantizer._calibrator) == MaxCalibrator)

    def test_quant_descriptors_are_piped_to_custom_quant_modules_if_has_kwargs(self):
        # ARRANGE
        class MyBlock(nn.Module):
            def __init__(self, in_feats, out_feats) -> None:
                super().__init__()
                self.in_feats = in_feats
                self.out_feats = out_feats
                self.flatten = nn.Flatten()
                self.linear = nn.Linear(in_feats, out_feats)

            def forward(self, x):
                return self.linear(self.flatten(x))

        class MyQuantizedBlock(SGQuantMixin):
            # NOTE: if **kwargs are existing, then quant descriptors are passed there!
            # NOTE: because we don't override `from_float`,
            #       then the float instance should have `in_feats` and `out_feats` as state
            def __init__(self, in_feats, out_feats, **kwargs) -> None:
                super().__init__()
                self.flatten = nn.Flatten()
                self.linear = quant_nn.QuantLinear(
                    in_feats,
                    out_feats,
                    quant_desc_input=kwargs["quant_desc_input"],
                    quant_desc_weight=kwargs["quant_desc_weight"],
                )

            def forward(self, x):
                return self.linear(self.flatten(x))

        class MyModel(nn.Module):
            def __init__(self, res, n_classes) -> None:
                super().__init__()
                self.conv = nn.Conv2d(3, 4, kernel_size=3, padding=1)
                self.my_block = QuantizedMapping(
                    float_module=MyBlock(4 * (res**2), n_classes),
                    quantized_target_class=MyQuantizedBlock,
                    input_quant_descriptor=QuantDescriptor(calib_method="max"),
                    weights_quant_descriptor=QuantDescriptor(calib_method="histogram"),
                )

            def forward(self, x):
                y = self.conv(x)
                return self.my_block(y)

        res = 32
        n_clss = 10
        module = MyModel(res, n_clss)

        # TEST
        q_util = SelectiveQuantizer()
        q_util.quantize_module(module)

        x = torch.rand(1, 3, res, res)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, n_clss))

        self.assertTrue(isinstance(module.conv, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(isinstance(module.my_block, MyQuantizedBlock))
        self.assertTrue(type(module.my_block.linear._input_quantizer._calibrator) == MaxCalibrator)
        self.assertTrue(type(module.my_block.linear._weight_quantizer._calibrator) == HistogramCalibrator)

    def test_quant_descriptors_are_piped_to_custom_quant_modules_if_expects_in_init(self):
        # ARRANGE
        class MyBlock(nn.Module):
            def __init__(self, in_feats, out_feats) -> None:
                super().__init__()
                self.in_feats = in_feats
                self.out_feats = out_feats
                self.flatten = nn.Flatten()
                self.linear = nn.Linear(in_feats, out_feats)

            def forward(self, x):
                return self.linear(self.flatten(x))

        class MyQuantizedBlock(SGQuantMixin):
            # NOTE: `since quant_desc_input`, `quant_desc_weight` are existing, then quant descriptors are passed there!
            # NOTE: because we don't override `from_float`,
            #       then the float instance should have `in_feats` and `out_feats` as state
            def __init__(self, in_feats, out_feats, quant_desc_input, quant_desc_weight) -> None:
                super().__init__()
                self.flatten = nn.Flatten()
                self.linear = quant_nn.QuantLinear(
                    in_feats,
                    out_feats,
                    quant_desc_input=quant_desc_input,
                    quant_desc_weight=quant_desc_weight,
                )

            def forward(self, x):
                return self.linear(self.flatten(x))

        class MyModel(nn.Module):
            def __init__(self, res, n_classes) -> None:
                super().__init__()
                self.conv = nn.Conv2d(3, 4, kernel_size=3, padding=1)
                self.my_block = QuantizedMapping(
                    float_module=MyBlock(4 * (res**2), n_classes),
                    quantized_target_class=MyQuantizedBlock,
                    input_quant_descriptor=QuantDescriptor(calib_method="max"),
                    weights_quant_descriptor=QuantDescriptor(calib_method="histogram"),
                )

            def forward(self, x):
                y = self.conv(x)
                return self.my_block(y)

        res = 32
        n_clss = 10
        module = MyModel(res, n_clss)

        # TEST
        q_util = SelectiveQuantizer()
        q_util.quantize_module(module)

        x = torch.rand(1, 3, res, res)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, n_clss))

        self.assertTrue(isinstance(module.conv, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(isinstance(module.my_block, MyQuantizedBlock))
        self.assertTrue(type(module.my_block.linear._input_quantizer._calibrator) == MaxCalibrator)
        self.assertTrue(type(module.my_block.linear._weight_quantizer._calibrator) == HistogramCalibrator)

    def test_quant_descriptors_are_not_piped_if_custom_quant_module_does_not_expect_them(self):
        # ARRANGE
        class MyBlock(nn.Module):
            def __init__(self, in_feats, out_feats) -> None:
                super().__init__()
                self.in_feats = in_feats
                self.out_feats = out_feats
                self.flatten = nn.Flatten()
                self.linear = nn.Linear(in_feats, out_feats)

            def forward(self, x):
                return self.linear(self.flatten(x))

        class MyQuantizedBlock(SGQuantMixin):
            # NOTE: because we don't override `from_float`,
            #       then the float instance should have `in_feats` and `out_feats` as state
            def __init__(self, in_feats, out_feats) -> None:
                super().__init__()
                self.flatten = nn.Flatten()
                self.linear = quant_nn.QuantLinear(in_feats, out_feats)

            def forward(self, x):
                return self.linear(self.flatten(x))

        class MyModel(nn.Module):
            def __init__(self, res, n_classes) -> None:
                super().__init__()
                self.conv = nn.Conv2d(3, 4, kernel_size=3, padding=1)
                self.my_block = QuantizedMapping(float_module=MyBlock(4 * (res**2), n_classes), quantized_target_class=MyQuantizedBlock)

            def forward(self, x):
                y = self.conv(x)
                return self.my_block(y)

        res = 32
        n_clss = 10
        module = MyModel(res, n_clss)

        # TEST
        q_util = SelectiveQuantizer()
        q_util.quantize_module(module)

        x = torch.rand(1, 3, res, res)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, n_clss))

        self.assertTrue(isinstance(module.conv, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(isinstance(module.my_block, MyQuantizedBlock))

    def test_custom_quantized_mappings_are_recursively_quantized_if_required(self):
        # ARRANGE
        class MyBlock(nn.Module):
            def __init__(self, in_feats, out_feats) -> None:
                super().__init__()
                self.in_feats = in_feats
                self.out_feats = out_feats
                self.flatten = nn.Flatten()
                self.linear = nn.Linear(in_feats, out_feats)

            def forward(self, x):
                return self.linear(self.flatten(x))

        class MyQuantizedBlock(SGQuantMixin):
            def __init__(self, in_feats, out_feats) -> None:
                super().__init__()
                self.flatten = nn.Flatten()
                self.linear = nn.Linear(in_feats, out_feats)

            def forward(self, x):
                return self.linear(self.flatten(x))

        class MyModel(nn.Module):
            def __init__(self, res, n_classes) -> None:
                super().__init__()
                self.conv = nn.Conv2d(3, 4, kernel_size=3, padding=1)
                self.my_block = QuantizedMapping(
                    float_module=MyBlock(4 * (res**2), n_classes),
                    quantized_target_class=MyQuantizedBlock,
                    action=QuantizedMetadata.ReplacementAction.REPLACE_AND_RECURE,
                )

            def forward(self, x):
                y = self.conv(x)
                return self.my_block(y)

        res = 32
        n_clss = 10
        module = MyModel(res, n_clss)

        # TEST
        q_util = SelectiveQuantizer()
        q_util.quantize_module(module)

        x = torch.rand(1, 3, res, res)

        # ASSERT
        with torch.no_grad():
            y = module(x)
            torch.testing.assert_close(y.size(), (1, n_clss))

        self.assertTrue(isinstance(module.conv, SelectiveQuantizer.mapping_instructions[nn.Conv2d].quantized_target_class))
        self.assertTrue(isinstance(module.my_block, MyQuantizedBlock))
        self.assertTrue(isinstance(module.my_block.linear, SelectiveQuantizer.mapping_instructions[nn.Linear].quantized_target_class))

    def test_torchvision_resnet_sg_vanilla_quantization_matches_pytorch_quantization(self):
        resnet_sg = torchvision.models.resnet50(pretrained=True)

        # SG SELECTIVE QUANTIZATION
        sq = SelectiveQuantizer(
            custom_mappings={
                torch.nn.Conv2d: QuantizedMetadata(
                    torch.nn.Conv2d,
                    quant_nn.QuantConv2d,
                    action=QuantizedMetadata.ReplacementAction.REPLACE,
                    input_quant_descriptor=QuantDescriptor(calib_method="histogram"),
                    weights_quant_descriptor=QuantDescriptor(calib_method="max", axis=0),
                ),
                torch.nn.Linear: QuantizedMetadata(
                    torch.nn.Linear,
                    quant_nn.QuantLinear,
                    action=QuantizedMetadata.ReplacementAction.REPLACE,
                    input_quant_descriptor=QuantDescriptor(calib_method="histogram"),
                    weights_quant_descriptor=QuantDescriptor(calib_method="max", axis=0),
                ),
                torch.nn.AdaptiveAvgPool2d: QuantizedMetadata(
                    torch.nn.AdaptiveAvgPool2d,
                    quant_nn.QuantAdaptiveAvgPool2d,
                    action=QuantizedMetadata.ReplacementAction.REPLACE,
                    input_quant_descriptor=QuantDescriptor(calib_method="max"),
                ),
            },
            default_per_channel_quant_modules=True,
        )

        sq.quantize_module(resnet_sg, preserve_state_dict=True)

        # PYTORCH-QUANTIZATION
        quant_desc_input = QuantDescriptor(calib_method="histogram")
        quant_nn.QuantConv2d.set_default_quant_desc_input(quant_desc_input)
        quant_nn.QuantLinear.set_default_quant_desc_input(quant_desc_input)

        quant_modules.initialize()
        resnet_pyquant = torchvision.models.resnet50(pretrained=True)
        quant_modules.deactivate()

        for (n1, p1), (n2, p2) in zip(resnet_sg.named_parameters(), resnet_pyquant.named_parameters()):
            torch.testing.assert_allclose(p1, p2)

        x = torch.rand(1, 3, 128, 128)
        with torch.no_grad():
            y_pyquant = resnet_pyquant(x)
            y_sg = resnet_sg(x)
            torch.testing.assert_close(y_sg, y_pyquant)

    def test_sg_resnet_sg_vanilla_quantization_matches_pytorch_quantization(self):

        # SG SELECTIVE QUANTIZATION
        sq = SelectiveQuantizer(
            custom_mappings={
                torch.nn.Conv2d: QuantizedMetadata(
                    torch.nn.Conv2d,
                    quant_nn.QuantConv2d,
                    action=QuantizedMetadata.ReplacementAction.REPLACE,
                    input_quant_descriptor=QuantDescriptor(calib_method="histogram"),
                    weights_quant_descriptor=QuantDescriptor(calib_method="max", axis=0),
                ),
                torch.nn.Linear: QuantizedMetadata(
                    torch.nn.Linear,
                    quant_nn.QuantLinear,
                    action=QuantizedMetadata.ReplacementAction.REPLACE,
                    input_quant_descriptor=QuantDescriptor(calib_method="histogram"),
                    weights_quant_descriptor=QuantDescriptor(calib_method="max", axis=0),
                ),
                torch.nn.AdaptiveAvgPool2d: QuantizedMetadata(
                    torch.nn.AdaptiveAvgPool2d,
                    quant_nn.QuantAdaptiveAvgPool2d,
                    action=QuantizedMetadata.ReplacementAction.REPLACE,
                    input_quant_descriptor=QuantDescriptor(calib_method="max"),
                ),
            },
            default_per_channel_quant_modules=True,
        )

        resnet_sg: nn.Module = super_gradients.training.models.get("resnet50", pretrained_weights="imagenet", num_classes=1000)
        sq.quantize_module(resnet_sg, preserve_state_dict=True)

        # PYTORCH-QUANTIZATION
        quant_desc_input = QuantDescriptor(calib_method="histogram")
        quant_nn.QuantConv2d.set_default_quant_desc_input(quant_desc_input)
        quant_nn.QuantLinear.set_default_quant_desc_input(quant_desc_input)

        quant_modules.initialize()
        resnet_pyquant: nn.Module = super_gradients.training.models.get("resnet50", pretrained_weights="imagenet", num_classes=1000)

        quant_modules.deactivate()

        for (n1, p1), (n2, p2) in zip(resnet_sg.named_parameters(), resnet_pyquant.named_parameters()):
            torch.testing.assert_allclose(p1, p2)

        x = torch.rand(1, 3, 128, 128)
        with torch.no_grad():
            y_pyquant = resnet_pyquant(x)
            y_sg = resnet_sg(x)
            torch.testing.assert_close(y_sg, y_pyquant)


if __name__ == "__main__":
    unittest.main()