Backend-agnostic telemetry update

openvinotoolkit · Oct 18, 2024 · 9263d35 · 9263d35
1 parent 7385c41
commit 9263d35
Show file tree

Hide file tree

Showing 7 changed files with 68 additions and 46 deletions.
diff --git a/nncf/onnx/quantization/quantize_model.py b/nncf/onnx/quantization/quantize_model.py
@@ -32,15 +32,11 @@
 from nncf.quantization.algorithms.post_training.algorithm import PostTrainingQuantization
 from nncf.quantization.quantize_model import quantize_with_tune_hyperparams
 from nncf.quantization.quantize_model import warning_model_no_batchwise_support
-from nncf.quantization.telemetry_extractors import CompressionStartedWithQuantizeApi
 from nncf.scopes import IgnoredScope
-from nncf.telemetry import tracked_function
-from nncf.telemetry.events import NNCF_ONNX_CATEGORY
 
 TTensor = TypeVar("TTensor")
 
 
-@tracked_function(NNCF_ONNX_CATEGORY, [CompressionStartedWithQuantizeApi(), "target_device", "preset"])
 def quantize_impl(
     model: onnx.ModelProto,
     calibration_dataset: Dataset,

diff --git a/nncf/openvino/quantization/quantize_model.py b/nncf/openvino/quantization/quantize_model.py
@@ -49,18 +49,12 @@
 from nncf.quantization.quantize_model import is_model_no_batchwise_support
 from nncf.quantization.quantize_model import quantize_with_tune_hyperparams
 from nncf.quantization.quantize_model import warning_model_no_batchwise_support
-from nncf.quantization.telemetry_extractors import CompressionStartedWithCompressWeightsApi
-from nncf.quantization.telemetry_extractors import CompressionStartedWithQuantizeApi
-from nncf.quantization.telemetry_extractors import CompressionStartedWithQuantizeWithAccuracyControlApi
 from nncf.scopes import IgnoredScope
 from nncf.scopes import validate_ignored_scope
-from nncf.telemetry.decorator import tracked_function
-from nncf.telemetry.events import NNCF_OV_CATEGORY
 
 TTensor = TypeVar("TTensor")
 
 
-@tracked_function(NNCF_OV_CATEGORY, [CompressionStartedWithQuantizeApi(), "target_device", "preset"])
 def native_quantize_if_op_impl(
     model: ov.Model,
     calibration_dataset: Dataset,
@@ -143,7 +137,6 @@ def _extract_all_subgraphs(model: ov.Model, current_id: str) -> None:
     return quantized_model
 
 
-@tracked_function(NNCF_OV_CATEGORY, [CompressionStartedWithQuantizeApi(), "target_device", "preset"])
 def native_quantize_impl(
     model: ov.Model,
     calibration_dataset: Dataset,
@@ -191,10 +184,6 @@ def native_quantize_impl(
     return quantized_model
 
 
-@tracked_function(
-    NNCF_OV_CATEGORY,
-    [CompressionStartedWithQuantizeWithAccuracyControlApi(), "target_device", "preset", "max_drop", "drop_type"],
-)
 def quantize_with_accuracy_control_impl(
     model: ov.Model,
     calibration_dataset: Dataset,
@@ -369,18 +358,6 @@ def quantize_impl(
     )
 
 
-@tracked_function(
-    NNCF_OV_CATEGORY,
-    [
-        CompressionStartedWithCompressWeightsApi(),
-        "mode",
-        "awq",
-        "scale_estimation",
-        "gptq",
-        "lora_correction",
-        "backup_mode",
-    ],
-)
 def compress_weights_impl(
     model: ov.Model,
     dataset: Dataset,

diff --git a/nncf/quantization/quantize_model.py b/nncf/quantization/quantize_model.py
@@ -36,7 +36,12 @@
 from nncf.quantization.algorithms.hyperparameter_tuner.algorithm import HyperparameterTuner
 from nncf.quantization.algorithms.hyperparameter_tuner.param_grid import get_quantization_param_grids
 from nncf.quantization.algorithms.post_training.pipeline import create_ptq_pipeline
+from nncf.quantization.telemetry_extractors import CompressionStartedWithCompressWeightsApi
+from nncf.quantization.telemetry_extractors import CompressionStartedWithQuantizeApi
+from nncf.quantization.telemetry_extractors import CompressionStartedWithQuantizeWithAccuracyControlApi
 from nncf.scopes import IgnoredScope
+from nncf.telemetry.decorator import tracked_function
+from nncf.telemetry.events import MODEL_BASED_CATEGORY
 
 TTensor = TypeVar("TTensor")
 
@@ -111,6 +116,14 @@ def _update_advanced_quantization_parameters(
 
 
 @api(canonical_alias="nncf.quantize")
+@tracked_function(
+    MODEL_BASED_CATEGORY,
+    [
+        CompressionStartedWithQuantizeApi(),
+        "target_device",
+        "preset",
+    ],
+)
 def quantize(
     model: TModel,
     calibration_dataset: Dataset,
@@ -265,6 +278,16 @@ def wrapper(*args, **kwargs):
 
 
 @api(canonical_alias="nncf.quantize_with_accuracy_control")
+@tracked_function(
+    MODEL_BASED_CATEGORY,
+    [
+        CompressionStartedWithQuantizeWithAccuracyControlApi(),
+        "target_device",
+        "preset",
+        "max_drop",
+        "drop_type",
+    ],
+)
 def quantize_with_accuracy_control(
     model: TModel,
     calibration_dataset: Dataset,
@@ -380,6 +403,18 @@ def quantize_with_accuracy_control(
 
 
 @api(canonical_alias="nncf.compress_weights")
+@tracked_function(
+    MODEL_BASED_CATEGORY,
+    [
+        CompressionStartedWithCompressWeightsApi(),
+        "mode",
+        "awq",
+        "scale_estimation",
+        "gptq",
+        "lora_correction",
+        "backup_mode",
+    ],
+)
 def compress_weights(
     model: TModel,
     mode=CompressWeightsMode.INT8_ASYM,

diff --git a/nncf/telemetry/decorator.py b/nncf/telemetry/decorator.py
@@ -12,7 +12,9 @@
 import inspect
 from typing import Callable, List, Union
 
+from nncf.telemetry.events import MODEL_BASED_CATEGORY
 from nncf.telemetry.events import get_current_category
+from nncf.telemetry.events import get_model_based_category
 from nncf.telemetry.events import telemetry_category
 from nncf.telemetry.extractors import CollectedEvent
 from nncf.telemetry.extractors import TelemetryExtractor
@@ -49,6 +51,11 @@ def __call__(self, fn: Callable) -> Callable:
         def wrapped(*args, **kwargs):
             bound_args = fn_signature.bind(*args, **kwargs)
             bound_args.apply_defaults()
+
+            if self._category == MODEL_BASED_CATEGORY:
+                model_argument = bound_args.arguments.get("model", None)
+                self._category = get_model_based_category(model_argument)
+
             events: List[CollectedEvent] = []
             for collector in self._collectors:
                 argname = collector.argname

diff --git a/nncf/telemetry/events.py b/nncf/telemetry/events.py
@@ -10,15 +10,24 @@
 # limitations under the License.
 
 from contextlib import contextmanager
-from typing import Optional
+from typing import Optional, TypeVar
 
+from nncf.common.utils.backend import BackendType
+from nncf.common.utils.backend import get_backend
+
+# Backend categories
 NNCF_TF_CATEGORY = "nncf_tf"
 NNCF_PT_CATEGORY = "nncf_pt"
 NNCF_ONNX_CATEGORY = "nncf_onnx"
 NNCF_OV_CATEGORY = "nncf_ov"
 
+# Dynamic categories
+MODEL_BASED_CATEGORY = "model_based"
+
 CURRENT_CATEGORY = None
 
+TModel = TypeVar("TModel")
+
 
 def _set_current_category(category: str):
     global CURRENT_CATEGORY
@@ -29,6 +38,21 @@ def get_current_category() -> Optional[str]:
     return CURRENT_CATEGORY
 
 
+def get_model_based_category(model: TModel) -> str:
+    category_by_backend = {
+        BackendType.ONNX: NNCF_ONNX_CATEGORY,
+        BackendType.OPENVINO: NNCF_OV_CATEGORY,
+        BackendType.TORCH: NNCF_PT_CATEGORY,
+        BackendType.TENSORFLOW: NNCF_TF_CATEGORY,
+    }
+    category = None
+    if model is not None:
+        model_backend = get_backend(model)
+        category = category_by_backend[model_backend]
+
+    return category
+
+
 @contextmanager
 def telemetry_category(category: str) -> str:
     previous_category = get_current_category()

diff --git a/nncf/telemetry/wrapper.py b/nncf/telemetry/wrapper.py
@@ -16,10 +16,10 @@
 from typing import Callable, Optional
 from unittest.mock import MagicMock
 
-from nncf import __version__
 from nncf.common.logging import nncf_logger
 from nncf.definitions import NNCF_CI_ENV_VAR_NAME
 from nncf.definitions import NNCF_DEV_ENV_VAR_NAME
+from nncf.version import __version__ as __version__
 
 NNCFTelemetryStub = MagicMock
 

diff --git a/nncf/torch/quantization/quantize_model.py b/nncf/torch/quantization/quantize_model.py
@@ -29,18 +29,13 @@
 from nncf.quantization.algorithms.post_training.algorithm import PostTrainingQuantization
 from nncf.quantization.algorithms.weight_compression.algorithm import WeightCompression
 from nncf.quantization.quantize_model import warning_model_no_batchwise_support
-from nncf.quantization.telemetry_extractors import CompressionStartedWithCompressWeightsApi
-from nncf.quantization.telemetry_extractors import CompressionStartedWithQuantizeApi
 from nncf.scopes import IgnoredScope
-from nncf.telemetry.decorator import tracked_function
-from nncf.telemetry.events import NNCF_PT_CATEGORY
 from nncf.torch.graph.operator_metatypes import OPERATIONS_OUTPUT_HAS_NO_BATCH_AXIS
 from nncf.torch.model_creation import wrap_model
 
 DEFAULT_RANGE_TYPE = "mean_min_max"
 
 
-@tracked_function(NNCF_PT_CATEGORY, [CompressionStartedWithQuantizeApi(), "target_device", "preset"])
 def quantize_impl(
     model: torch.nn.Module,
     calibration_dataset: Dataset,
@@ -86,18 +81,6 @@ def quantize_impl(
     return quantized_model
 
 
-@tracked_function(
-    NNCF_PT_CATEGORY,
-    [
-        CompressionStartedWithCompressWeightsApi(),
-        "mode",
-        "awq",
-        "scale_estimation",
-        "gptq",
-        "lora_correction",
-        "backup_mode",
-    ],
-)
 def compress_weights_impl(
     model: torch.nn.Module,
     dataset: Dataset,