Feature/quantizer factory refactoring

documentation/source/ptq_qat.md

@@ -259,17 +259,14 @@ If you prefer more of a DIY approach, here is the code sample:
 
 ```python
 import torch
-from super_gradients.training.utils.quantization.export import export_quantized_module_to_onnx
-
+from super_gradients.conversion.onnx import export_to_onnx
 onnx_filename = f"qat_model_1x3x224x224.onnx"
 
 dummy_input = torch.randn([1, 3, 224, 224], device="cpu")
-export_quantized_module_to_onnx(
-    model=quantized_model.cpu(),
+export_to_onnx(
+    model=quantized_model.eval(),
+    model_input=dummy_input,
     onnx_filename=onnx_filename,
-    input_shape=[1, 3, 224, 224],
-    input_size=[1, 3, 224, 224],
-    train=False,
 )
 ```
 

src/super_gradients/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "3.7.1"
+__version__ = "3.8.0"
 
 from super_gradients.common import init_trainer, is_distributed, object_names
 from super_gradients.training import losses, utils, datasets_utils, DataAugmentation, Trainer, KDTrainer, QATTrainer

src/super_gradients/common/factories/exporter_factory.py

@@ -0,0 +1,7 @@
+from super_gradients.common.factories.base_factory import BaseFactory
+from super_gradients.common.registry.registry import ALL_EXPORTERS
+
+
+class ExporterFactory(BaseFactory):
+    def __init__(self):
+        super().__init__(ALL_EXPORTERS)

src/super_gradients/common/factories/quantizer_factory.py

@@ -0,0 +1,7 @@
+from super_gradients.common.factories.base_factory import BaseFactory
+from super_gradients.common.registry.registry import ALL_QUANTIZER
+
+
+class QuantizerFactory(BaseFactory):
+    def __init__(self):
+        super().__init__(ALL_QUANTIZER)

src/super_gradients/common/registry/registry.py

@@ -194,3 +194,9 @@ def warn_if_deprecated(name: str, registry: dict):
 
 PROCESSINGS = {}
 register_processing = create_register_decorator(registry=PROCESSINGS)
+
+ALL_QUANTIZER = {}
+register_quantizer = create_register_decorator(registry=ALL_QUANTIZER)
+
+ALL_EXPORTERS = {}
+register_exporter = create_register_decorator(registry=ALL_EXPORTERS)

src/super_gradients/conversion/abstract_exporter.py

@@ -0,0 +1,20 @@
+import abc
+
+
+class AbstractExporter(abc.ABC):
+    """
+    An abstract class for exporting a model from ONNX representation to a specific framework.
+    For instance, ONNX model can be exported to TFLite, OpenVINO or CoreML formats.
+    This can be done by subclassing from this class and implementing the `export_from_onnx` method.
+    """
+
+    @abc.abstractmethod
+    def export_from_onnx(self, source_onnx: str, output_file: str) -> str:
+        """
+        Exports a model from ONNX representation to an output file.
+        A output filename extension and it's content should be determined by the subclass.
+        :param source_onnx: Input ONNX model file path.
+        :param output_file: Output file path of the exported model.
+        :return: Output file path of the exported model.
+        """
+        pass

src/super_gradients/conversion/export_params.py

@@ -1,7 +1,7 @@
 import dataclasses
 from typing import Optional, Tuple
 
-from super_gradients.conversion.conversion_enums import ExportTargetBackend, DetectionOutputFormatMode
+from super_gradients.conversion.conversion_enums import DetectionOutputFormatMode
 
 
 @dataclasses.dataclass
@@ -10,10 +10,6 @@ class ExportParams:
     Parameters for exporting a model to ONNX format in PTQ/QAT methods of Trainer.
     Most of the parameters are related ot ExportableObjectDetectionModel.export method.
 
-    :param output_onnx_path: The path to save the ONNX model.
-           If None, the ONNX filename will use current experiment dir folder
-           and the output filename will reflect model input shape & whether it's a PTQ or QAT model.
-
     :param batch_size: The batch size for the ONNX model. Default is 1.
 
     :param input_image_shape: The input image shape (rows, cols) for the ONNX model.
@@ -58,8 +54,6 @@ class ExportParams:
            Relevant only for object detection models and only if postprocessing is True.
     """
 
-    output_onnx_path: Optional[str] = None
-    engine: Optional[ExportTargetBackend] = None
     batch_size: int = 1
     input_image_shape: Optional[Tuple[int, int]] = None
     preprocessing: bool = True
@@ -74,3 +68,4 @@ class ExportParams:
     detection_max_predictions_per_image: Optional[int] = None
     detection_predictions_format: DetectionOutputFormatMode = DetectionOutputFormatMode.BATCH_FORMAT
     detection_num_pre_nms_predictions: int = 1000
+    detection_postprocessing_use_tensorrt_nms: bool = False

src/super_gradients/conversion/onnx/__init__.py

@@ -0,0 +1,3 @@
+from .export_to_onnx import export_to_onnx
+
+__all__ = ["export_to_onnx"]

src/super_gradients/conversion/onnx_exporter.py

@@ -0,0 +1,17 @@
+import shutil
+
+from super_gradients.common.abstractions.abstract_logger import get_logger
+from super_gradients.common.registry.registry import register_exporter
+from super_gradients.conversion.abstract_exporter import AbstractExporter
+
+logger = get_logger(__name__)
+
+
+@register_exporter()
+class ONNXExporter(AbstractExporter):
+    def __init__(self):
+        pass
+
+    def export_from_onnx(self, source_onnx: str, output_file: str) -> str:
+        shutil.copy(source_onnx, output_file)
+        return output_file