Update check_nncf_graph.

Add reference graphs for anomaly models.

external/anomaly/tasks/inference.py

@@ -300,8 +300,13 @@ def save_model(self, output_model: ModelEntity) -> None:
     def _set_metadata(self, output_model: ModelEntity):
         output_model.set_data("image_threshold", self.model.image_threshold.value.cpu().numpy().tobytes())
         output_model.set_data("pixel_threshold", self.model.pixel_threshold.value.cpu().numpy().tobytes())
-        output_model.set_data("min", self.model.normalization_metrics.state_dict()["min"].cpu().numpy().tobytes())
-        output_model.set_data("max", self.model.normalization_metrics.state_dict()["max"].cpu().numpy().tobytes())
+        if hasattr(self.model, "normalization_metrics"):
+            output_model.set_data("min", self.model.normalization_metrics.state_dict()["min"].cpu().numpy().tobytes())
+            output_model.set_data("max", self.model.normalization_metrics.state_dict()["max"].cpu().numpy().tobytes())
+        else:
+            logger.warning(
+                "The model was not trained before saving. This will lead to incorrect normalization of the heatmaps."
+            )
 
     @staticmethod
     def _is_docker() -> bool:

external/anomaly/tests/anomaly_common.py

@@ -18,6 +18,10 @@
 from collections import namedtuple
 from copy import deepcopy
 from typing import List, Type
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from nncf.torch.nncf_network import NNCFNetwork
 
 from adapters.anomalib.data.mvtec import OteMvtecDataset
 from ote_sdk.entities.datasets import DatasetEntity
@@ -38,6 +42,7 @@
     OTETestTrainingEvaluationAction,
 )
 from ote_sdk.test_suite.training_tests_common import ROOT_PATH_KEY, make_paths_be_abs
+from tasks import NNCFTask
 
 logger = logging.getLogger(__name__)
 
@@ -107,3 +112,18 @@ def get_anomaly_domain_test_action_classes(
         OTETestNNCFExportEvaluationAction,
         OTETestNNCFGraphAction,
     ]
+
+
+def get_dummy_compressed_model(task: NNCFTask) -> "NNCFNetwork":
+    """
+    Return compressed model without initialization
+    """
+    from anomalib.utils.callbacks.nncf.utils import wrap_nncf_model
+
+    # Disable quantaizers initialization
+    for compression in task.optimization_config["nncf_config"]["compression"]:
+        if compression["algorithm"] == "quantization":
+            compression["initializer"] = {"batchnorm_adaptation": {"num_bn_adaptation_samples": 0}}
+
+    _, compressed_model = wrap_nncf_model(task.model, task.optimization_config["nncf_config"])
+    return compressed_model

external/anomaly/tests/test_configurable_parameters.py

@@ -21,8 +21,8 @@
 
 import pytest
 from adapters.anomalib.config import get_anomalib_config
-from configs.anomaly_classification.padim import PadimAnomalyClassificationConfig
-from configs.anomaly_classification.stfpm import STFPMAnomalyClassificationConfig
+from configs.classification.padim import PadimAnomalyClassificationConfig
+from configs.classification.stfpm import STFPMAnomalyClassificationConfig
 from ote_sdk.configuration.helper import convert, create
 
 from tests.helpers.config import get_config_and_task_name

external/anomaly/tests/test_ote_anomaly_classification_training.py

@@ -51,6 +51,7 @@
     _create_anomaly_dataset_and_labels_schema,
     _get_dataset_params_from_dataset_definitions,
     get_anomaly_domain_test_action_classes,
+    get_dummy_compressed_model,
 )
 
 logger = get_logger(__name__)
@@ -287,7 +288,7 @@ def _nncf_graph_params_factory() -> Dict:
                 "labels_schema": labels_schema,
                 "template_path": template_path,
                 "reference_dir": ote_current_reference_dir_fx,
-                "fn_get_compressed_model": None,  # NNCF not yet implemented in Anomaly
+                "fn_get_compressed_model": get_dummy_compressed_model,
             }
 
         params_factories_for_test_actions = {

external/anomaly/tests/test_ote_anomaly_segmentation_training.py

@@ -53,6 +53,7 @@
     _create_anomaly_dataset_and_labels_schema,
     _get_dataset_params_from_dataset_definitions,
     get_anomaly_domain_test_action_classes,
+    get_dummy_compressed_model,
 )
 
 logger = logging.getLogger(__name__)
@@ -433,7 +434,7 @@ def _nncf_graph_params_factory() -> Dict[str, Callable[[], Dict]]:
                 "labels_schema": labels_schema,
                 "template_path": template_path,
                 "reference_dir": ote_current_reference_dir_fx,
-                "fn_get_compressed_model": None,  # NNCF not yet implemented in Anomaly
+                "fn_get_compressed_model": get_dummy_compressed_model,
             }
 
         params_factories_for_test_actions = {

ote_sdk/ote_sdk/test_suite/fixtures.py

@@ -69,8 +69,6 @@ def ote_current_reference_dir_fx(ote_reference_root_dir_fx, current_test_paramet
     path = os.path.join(
         ote_reference_root_dir_fx, current_test_parameters_fx["model_name"]
     )
-    if not os.path.isdir(path):
-        return None
     return path
 
 

ote_sdk/ote_sdk/test_suite/training_tests_actions.py

@@ -9,7 +9,10 @@
 from abc import ABC, abstractmethod
 from collections import OrderedDict
 from copy import deepcopy
-from typing import List, Optional, Type
+from typing import TYPE_CHECKING, List, Optional, Type
+
+if TYPE_CHECKING:
+    from nncf.torch.nncf_network import NNCFNetwork
 
 import pytest
 
@@ -562,15 +565,28 @@ def __call__(self, data_collector: DataCollector, results_prev_stages: OrderedDi
 
 
 # TODO: think about move to special file
-def check_nncf_model_graph(model, path_to_dot):
+def check_nncf_model_graph(
+    model: "NNCFNetwork", path_to_dot: str, force_regen_dot: bool = False
+):
+    """
+    Compare graph of compressed model with reference.
+    """
     import networkx as nx
 
-    logger.info(f"Reference graph: {path_to_dot}")
-    load_graph = nx.drawing.nx_pydot.read_dot(path_to_dot)
-
     graph = model.get_graph()
     nx_graph = graph.get_graph_for_structure_analysis()
 
+    if force_regen_dot:
+        # Generate and rewrite reference graph
+        dir_path = os.path.dirname(path_to_dot)
+        if not os.path.exists(dir_path):
+            os.makedirs(dir_path)
+        nx.drawing.nx_pydot.write_dot(nx_graph, path_to_dot)
+        logger.warning(f"Reference graph was generated: {path_to_dot}")
+
+    logger.info(f"Reference graph: {path_to_dot}")
+    load_graph = nx.drawing.nx_pydot.read_dot(path_to_dot)
+
     for _, node in nx_graph.nodes(data=True):
         if "scope" in node:
             node.pop("scope")
@@ -592,6 +608,7 @@ def check_nncf_model_graph(model, path_to_dot):
 
 class OTETestNNCFGraphAction(BaseOTETestAction):
     _name = "nncf_graph"
+    _VAR_REGEN_DOT = "NNCF_TEST_REGEN_DOT"
 
     def __init__(
         self,
@@ -619,8 +636,13 @@ def _run_ote_nncf_graph(self, data_collector):
         if not is_nncf_enabled():
             pytest.skip("NNCF is not installed")
 
-        if not os.path.exists(self.reference_dir):
-            pytest.skip("Reference directory does not exist")
+        force_regen_dot = os.getenv(self._VAR_REGEN_DOT) is not None
+        if not force_regen_dot:
+            if self.reference_dir is None or not os.path.exists(self.reference_dir):
+                pytest.skip(
+                    f"Reference directory does not exist: {self.reference_dir}.\n"
+                    f"To generate reference graph set the global variable {self._VAR_REGEN_DOT}."
+                )
 
         params = ote_sdk_configuration_helper_create(
             model_template.hyper_parameters.data
@@ -650,16 +672,22 @@ def _run_ote_nncf_graph(self, data_collector):
         nncf_task_cls = get_impl_class(nncf_task_class_impl_path)
         nncf_task = nncf_task_cls(task_environment=environment_for_nncf)
 
-        path_to_ref_dot = os.path.join(
-            self.reference_dir, "nncf", f"{nncf_task._nncf_preset}.dot"
+        nncf_preset = (
+            nncf_task.nncf_preset
+            if hasattr(nncf_task, "nncf_preset")
+            else nncf_task._nncf_preset
         )
-        if not os.path.exists(path_to_ref_dot):
-            pytest.skip("Reference file does not exist: {}".format(path_to_ref_dot))
+        path_to_ref_dot = os.path.join(self.reference_dir, "nncf", f"{nncf_preset}.dot")
+        if not os.path.exists(path_to_ref_dot) and not force_regen_dot:
+            pytest.skip(
+                f"Reference file does not exist: {path_to_ref_dot}.\n"
+                f"To generate reference graph set the global variable {self._VAR_REGEN_DOT}."
+            )
 
         compressed_model = self.fn_get_compressed_model(nncf_task)
 
         assert check_nncf_model_graph(
-            compressed_model, path_to_ref_dot
+            compressed_model, path_to_ref_dot, force_regen_dot
         ), "Compressed model differs from the reference"
 
     def __call__(self, data_collector: DataCollector, results_prev_stages: OrderedDict):