🔨 Fix nncf key issue in nightly job

openvinotoolkit · May 5, 2022 · 60d9c12 · 60d9c12
1 parent edf9393
commit 60d9c12
Show file tree

Hide file tree

Showing 10 changed files with 20 additions and 14 deletions.
diff --git a/anomalib/models/cflow/config.yaml b/anomalib/models/cflow/config.yaml
@@ -86,7 +86,7 @@ trainer:
   move_metrics_to_cpu: false
   multiple_trainloader_mode: max_size_cycle
   num_nodes: 1
-  num_processes: 1
+  num_processes: null
   num_sanity_val_steps: 0
   overfit_batches: 0.0
   plugins: null

diff --git a/anomalib/models/dfkde/config.yaml b/anomalib/models/dfkde/config.yaml
@@ -70,7 +70,7 @@ trainer:
   move_metrics_to_cpu: false
   multiple_trainloader_mode: max_size_cycle
   num_nodes: 1
-  num_processes: 1
+  num_processes: null
   num_sanity_val_steps: 0
   overfit_batches: 0.0
   plugins: null

diff --git a/anomalib/models/dfm/config.yaml b/anomalib/models/dfm/config.yaml
@@ -74,7 +74,7 @@ trainer:
   move_metrics_to_cpu: false
   multiple_trainloader_mode: max_size_cycle
   num_nodes: 1
-  num_processes: 1
+  num_processes: null
   num_sanity_val_steps: 0
   overfit_batches: 0.0
   plugins: null

diff --git a/anomalib/models/ganomaly/config.yaml b/anomalib/models/ganomaly/config.yaml
@@ -92,7 +92,7 @@ trainer:
   move_metrics_to_cpu: false
   multiple_trainloader_mode: max_size_cycle
   num_nodes: 1
-  num_processes: 1
+  num_processes: null
   num_sanity_val_steps: 0
   overfit_batches: 0.0
   plugins: null

diff --git a/anomalib/models/padim/config.yaml b/anomalib/models/padim/config.yaml
@@ -84,7 +84,7 @@ trainer:
   move_metrics_to_cpu: false
   multiple_trainloader_mode: max_size_cycle
   num_nodes: 1
-  num_processes: 1
+  num_processes: null
   num_sanity_val_steps: 0
   overfit_batches: 0.0
   plugins: null

diff --git a/anomalib/models/patchcore/config.yaml b/anomalib/models/patchcore/config.yaml
@@ -83,7 +83,7 @@ trainer:
   move_metrics_to_cpu: false
   multiple_trainloader_mode: max_size_cycle
   num_nodes: 1
-  num_processes: 1
+  num_processes: null
   num_sanity_val_steps: 0
   overfit_batches: 0.0
   plugins: null

diff --git a/anomalib/models/stfpm/config.yaml b/anomalib/models/stfpm/config.yaml
@@ -93,7 +93,7 @@ trainer:
   move_metrics_to_cpu: false
   multiple_trainloader_mode: max_size_cycle
   num_nodes: 1
-  num_processes: 1
+  num_processes: null
   num_sanity_val_steps: 0
   overfit_batches: 0.0
   plugins: null

diff --git a/anomalib/utils/callbacks/nncf/callback.py b/anomalib/utils/callbacks/nncf/callback.py
@@ -38,9 +38,9 @@ class NNCFCallback(Callback):
                           If None model will not be exported.
     """
 
-    def __init__(self, nncf_config: Dict, export_dir: str = None):
+    def __init__(self, config: Dict, export_dir: str = None):
         self.export_dir = export_dir
-        self.nncf_config = NNCFConfig(nncf_config)
+        self.config = NNCFConfig(config)
         self.nncf_ctrl: Optional[CompressionAlgorithmController] = None
 
     # pylint: disable=unused-argument
@@ -54,10 +54,10 @@ def setup(self, trainer: pl.Trainer, pl_module: pl.LightningModule, stage: Optio
             return
 
         init_loader = InitLoader(trainer.datamodule.train_dataloader())  # type: ignore
-        nncf_config = register_default_init_args(self.nncf_config, init_loader)
+        config = register_default_init_args(self.config, init_loader)
 
         self.nncf_ctrl, pl_module.model = wrap_nncf_model(
-            model=pl_module.model, config=nncf_config, dataloader=trainer.datamodule.train_dataloader()  # type: ignore
+            model=pl_module.model, config=config, dataloader=trainer.datamodule.train_dataloader()  # type: ignore
         )
 
     def on_train_batch_start(

diff --git a/tests/helpers/model.py b/tests/helpers/model.py
@@ -64,7 +64,13 @@ def setup_model_train(
     config.dataset.category = category
     config.dataset.path = dataset_path
     config.project.log_images_to = []
-    config.trainer.gpus = device
+    config.trainer.devices = device
+    config.trainer.accelerator = "gpu" if device != 0 else "cpu"
+
+    # Remove legacy flags
+    for legacy_device in ["num_processes", "gpus", "ipus", "tpu_cores"]:
+        if legacy_device in config.trainer:
+            config.trainer[legacy_device] = None
 
     # If weight file is empty, remove the key from config
     if "weight_file" in config.model.keys() and weight_file == "":
@@ -73,7 +79,7 @@ def setup_model_train(
         config.model.weight_file = weight_file if not fast_run else "weights/last.ckpt"
 
     if nncf:
-        config.optimization.nncf.apply = True
+        config.optimization["nncf"] = {"apply": True, "input_info": {"sample_size": None}}
         config = update_nncf_config(config)
         config.init_weights = None
 

diff --git a/tests/nightly/models/test_model_nightly.py b/tests/nightly/models/test_model_nightly.py
@@ -88,7 +88,7 @@ def _test_metrics(self, trainer, config, model, datamodule):
         thresholds = OmegaConf.load("tests/nightly/models/performance_thresholds.yaml")
 
         threshold = thresholds[config.model.name][config.dataset.category]
-        if "optimization" in config.keys() and config.optimization.nncf.apply:
+        if "optimization" in config.keys() and "nncf" in config.optimization.keys() and config.optimization.nncf.apply:
             threshold = threshold.nncf
         if not (
             np.isclose(results["image_AUROC"], threshold["image_AUROC"], rtol=0.02)