Import backend modules in nncf/__init__.py

README.md

@@ -242,7 +242,6 @@ Here is an example of Accuracy Aware Quantization pipeline where model weights a
 
 ```python
 import nncf
-import nncf.torch
 import torch
 from torchvision import datasets, models
 
@@ -296,7 +295,7 @@ Here is an example of Accuracy Aware RB Sparsification pipeline where model weig
 
 ```python
 import torch
-import nncf.torch  # Important - must be imported before any other external package that depends on torch
+import nncf  # Important - must be imported before any other external package that depends on torch
 
 from nncf import NNCFConfig
 from nncf.torch import create_compressed_model, register_default_init_args

docs/usage/training_time_compression/other_algorithms/Usage.md

@@ -22,7 +22,7 @@ The framework is designed so that the modifications to your original training co
 
     ```python
     import torch
-    import nncf.torch  # Important - must be imported before any other external package that depends on torch
+    import nncf  # Important - must be imported before any other external package that depends on torch
     from nncf import NNCFConfig, create_compressed_model, load_state
     ```
 

docs/usage/training_time_compression/quantization_aware_training/Usage.md

@@ -81,7 +81,7 @@ The quantized model saving allows to load quantized modules to the target model
 requires only example input for the target module, corresponding NNCF config and the quantized model state dict.
 
 ```python
-import nncf.torch
+import nncf
 
 # save part
 quantized_model = nncf.quantize(model, calibration_dataset)

examples/llm_compression/torch/qat_with_lora/main.py

@@ -31,7 +31,6 @@
 from whowhatbench import TextEvaluator
 
 import nncf
-import nncf.torch
 from nncf.common.logging.track_progress import track
 from nncf.data.dataset import Dataset
 from nncf.parameters import CompressionFormat

examples/post_training_quantization/torch_fx/resnet18/main.py

@@ -29,7 +29,6 @@
 from torch._dynamo.exc import BackendCompilerFailed
 
 import nncf
-import nncf.torch
 from nncf.common.utils.helpers import create_table
 from nncf.common.utils.os import is_windows
 

examples/quantization_aware_training/torch/resnet18/main.py

@@ -31,7 +31,6 @@
 from torch.jit import TracerWarning
 
 import nncf
-import nncf.torch
 from nncf.common.utils.helpers import create_table
 
 warnings.filterwarnings("ignore", category=TracerWarning)

nncf/__init__.py

@@ -90,3 +90,15 @@
         "Please install one of the supported frameworks above in order to use NNCF on top of it.\n"
         "See the installation guide at https://github.com/openvinotoolkit/nncf#installation-guide for help."
     )
+
+if _AVAILABLE_FRAMEWORKS["torch"]:
+    from nncf import torch as torch
+
+if _AVAILABLE_FRAMEWORKS["tensorflow"]:
+    from nncf import tensorflow as tensorflow
+
+if _AVAILABLE_FRAMEWORKS["onnx"]:
+    from nncf import onnx as onnx
+
+if _AVAILABLE_FRAMEWORKS["openvino"]:
+    from nncf import openvino as openvino

nncf/experimental/torch/fx/transformations.py

@@ -19,7 +19,6 @@
 from torch.quantization.fake_quantize import FakeQuantize
 
 import nncf
-import nncf.torch
 from nncf.common.graph.graph import NNCFNode
 from nncf.common.graph.transformations.commands import TargetType
 from nncf.experimental.torch.fx.constant_folding import constant_fold

nncf/torch/dynamic_graph/patch_pytorch.py

@@ -283,6 +283,8 @@ def __init__(self, name: str, namespace, op):
 ORIGINAL_OPERATORS: List[OriginalOpInfo] = []
 ORIGINAL_CALL = torch.nn.Module.__call__
 _ORIG_JIT_SCRIPT = None
+_ORIG_JIT_SCRIPT_IF_TRACE = None
+
 _ORIG_JIT_TRACE_MAKE_MODULE = None
 _ORIG_TORCH_COMPILE: Union[Callable, None] = None
 
@@ -312,6 +314,8 @@ def patch_torch_jit():
 
     # Patch torch.jit._script_if_tracing because it references an original
     # unpatched torch.jit.script and the patching above does not affect it
+    global _ORIG_JIT_SCRIPT_IF_TRACE
+    _ORIG_JIT_SCRIPT_IF_TRACE = getattr(torch.jit, "_script_if_tracing")
     setattr(torch.jit, "_script_if_tracing", torch_jit_script_if_tracing)
 
 

tests/torch/conftest.py

@@ -9,6 +9,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
+
+# Enable patching of torch functions
+# Should be set before any import of nncf
+os.environ["NNCF_TORCH_LEGACY_TRACING"] = "1"
+
 import random
 from pathlib import Path
 
@@ -122,9 +127,6 @@ def pytest_configure(config: Config):
         if config.getoption(f"--regen-{regen_option}", False):
             os.environ[f"NNCF_TEST_REGEN_{regen_option.upper()}"] = "1"
 
-    # Enable patching of torch functions
-    os.environ["NNCF_TORCH_LEGACY_TRACING"] = "1"
-
 
 @pytest.fixture(scope="module")
 def dataset_dir(request: FixtureRequest):

tests/torch/pytorch_patch_isolated.py

@@ -40,9 +40,9 @@ def clean_source_code(code_source):
 @pytest.mark.skipif(ISOLATION_RUN_ENV_VAR not in os.environ, reason="Should be run via isolation proxy")
 def test_jit_if_tracing_script_source_equals():
     # Get original torch.jit._script_if_tracing source
-    torch_source = clean_source_code(inspect.getsource(torch.jit._script_if_tracing))
+    from nncf.torch.dynamic_graph.patch_pytorch import _ORIG_JIT_SCRIPT_IF_TRACE
 
-    import nncf.torch  # noqa: F401
+    torch_source = clean_source_code(inspect.getsource(_ORIG_JIT_SCRIPT_IF_TRACE))
 
     # Get torch.jit._script_if_tracing source after patching was performed
     nncf_source = clean_source_code(inspect.getsource(torch.jit._script_if_tracing))
@@ -109,7 +109,7 @@ def forward(self, x):
 def test_compile():
     compile_forward = os.environ.get("COMPILE_FORWARD", None) == "1"
     before_nncf = compile_and_run_test_model(compile_forward)
-    import nncf.torch  # noqa: F401
+    import nncf  # noqa: F401
 
     after_nncf = compile_and_run_test_model(compile_forward)
     assert torch.allclose(before_nncf, after_nncf)

tests/torch2/function_hook/quantization/strip/test_strip_dequantize.py

@@ -18,7 +18,6 @@
 from torch import nn
 
 import nncf
-import nncf.torch
 from nncf.common.quantization.structs import QuantizationScheme
 from nncf.parameters import CompressWeightsMode
 from nncf.parameters import StripFormat

tests/torch2/function_hook/quantization/strip/test_strip_native.py

@@ -17,7 +17,6 @@
 from torch.quantization import FakeQuantize
 
 import nncf
-import nncf.torch
 from nncf.quantization.advanced_parameters import OverflowFix
 from nncf.torch.function_hook.wrapper import get_hook_storage
 from tests.common.quantization.data_generators import generate_lazy_sweep_data

tests/torch2/test_patching.py

@@ -18,7 +18,6 @@
 
 def test_patching():
     # Check that patching torch functions is disabled
-    import nncf.torch  # noqa: F401
 
     with pytest.raises(AttributeError):
         getattr(torch.relu, "_original_op")