Quantization-aware training with NNCF

src/anomalib/callbacks/__init__.py

@@ -5,13 +5,10 @@
 
 
 import logging
-from importlib import import_module
-from pathlib import Path
 
-import yaml
 from jsonargparse import Namespace
 from lightning.pytorch.callbacks import Callback
-from omegaconf import DictConfig, ListConfig, OmegaConf
+from omegaconf import DictConfig, ListConfig
 
 from .checkpoint import ModelCheckpoint
 from .graph import GraphLogger
@@ -48,17 +45,4 @@ def get_callbacks(config: DictConfig | ListConfig | Namespace) -> list[Callback]
         load_model = LoadModelCallback(config.ckpt_path)
         callbacks.append(load_model)
 
-    if "optimization" in config and "nncf" in config.optimization and config.optimization.nncf.apply:
-        # NNCF wraps torch's jit which conflicts with kornia's jit calls.
-        # Hence, nncf is imported only when required
-        nncf_module = import_module("anomalib.utils.callbacks.nncf.callback")
-        nncf_callback = nncf_module.NNCFCallback
-        nncf_config = yaml.safe_load(OmegaConf.to_yaml(config.optimization.nncf))
-        callbacks.append(
-            nncf_callback(
-                config=nncf_config,
-                export_dir=str(Path(config.project.path) / "compressed"),
-            ),
-        )
-
     return callbacks