nncf.build_graph

src/nncf/__init__.py

@@ -12,6 +12,7 @@
 Neural Network Compression Framework (NNCF) for enhanced OpenVINO™ inference.
 """
 
+from nncf.common.factory import build_graph as build_graph
 from nncf.common.logging import nncf_logger as nncf_logger
 from nncf.common.logging.logger import disable_logging as disable_logging
 from nncf.common.logging.logger import set_log_level as set_log_level

src/nncf/common/factory.py

@@ -8,59 +8,71 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from __future__ import annotations
 
-from typing import Any, TypeVar, cast
+from typing import TYPE_CHECKING, Any, TypeVar, cast
 
 import nncf
-from nncf.common.engine import Engine
-from nncf.common.graph.graph import NNCFGraph
-from nncf.common.graph.model_transformer import ModelTransformer
-from nncf.common.graph.transformations.command_creation import CommandCreator
-from nncf.common.tensor_statistics import aggregator
+from nncf.common.utils.api_marker import api
 from nncf.common.utils.backend import BackendType
 from nncf.common.utils.backend import get_backend
-from nncf.data.dataset import Dataset
+
+if TYPE_CHECKING:
+    from nncf.common.engine import Engine
+    from nncf.common.graph.graph import NNCFGraph
+    from nncf.common.graph.model_transformer import ModelTransformer
+    from nncf.common.graph.transformations.command_creation import CommandCreator
+    from nncf.common.tensor_statistics.aggregator import StatisticsAggregator
+    from nncf.data.dataset import Dataset
 
 TModel = TypeVar("TModel")
 
 
-class NNCFGraphFactory:
-    @staticmethod
-    def create(model: TModel) -> NNCFGraph:
-        """
-        Factory method to create backend-specific NNCFGraph instance based on the input model.
+@api(canonical_alias="nncf.build_graph")
+def build_graph(model: TModel, *, example_input: Any = None) -> NNCFGraph:
+    """
+    Builds NNCFGraph from the given model.
 
-        :param model: backend-specific model instance
-        :return: backend-specific NNCFGraph instance
-        """
-        model_backend = get_backend(model)
-        if model_backend == BackendType.ONNX:
-            from onnx import ModelProto  # type: ignore
+    :param model: The model from which to build the graph.
+    :param example_input: Example input to the model, required for some backends.
+    :return: The constructed NNCFGraph.
+    """
+    model_backend = get_backend(model)
+    if model_backend == BackendType.ONNX:
+        from onnx import ModelProto  # type: ignore
 
-            from nncf.onnx.graph.nncf_graph_builder import GraphConverter as ONNXGraphConverter
+        from nncf.onnx.graph.nncf_graph_builder import GraphConverter as ONNXGraphConverter
 
-            return ONNXGraphConverter.create_nncf_graph(cast(ModelProto, model))
-        if model_backend == BackendType.OPENVINO:
-            from openvino import Model  # type: ignore
+        return ONNXGraphConverter.create_nncf_graph(cast(ModelProto, model))
+    if model_backend == BackendType.OPENVINO:
+        from openvino import Model  # type: ignore
 
-            from nncf.openvino.graph.nncf_graph_builder import GraphConverter as OVGraphConverter
+        from nncf.openvino.graph.nncf_graph_builder import GraphConverter as OVGraphConverter
 
-            return OVGraphConverter.create_nncf_graph(cast(Model, model))
-        if model_backend == BackendType.TORCH_FX:
-            from torch.fx import GraphModule
+        return OVGraphConverter.create_nncf_graph(cast(Model, model))
+    if model_backend == BackendType.TORCH_FX:
+        from torch.fx import GraphModule
 
-            from nncf.experimental.torch.fx.nncf_graph_builder import GraphConverter as FXGraphConverter
+        from nncf.experimental.torch.fx.nncf_graph_builder import GraphConverter as FXGraphConverter
 
-            return FXGraphConverter.create_nncf_graph(cast(GraphModule, model))
-        if model_backend == BackendType.TORCH:
-            from nncf.torch.function_hook.nncf_graph.nncf_graph_builder import GraphModelWrapper
+        return FXGraphConverter.create_nncf_graph(cast(GraphModule, model))
+    if model_backend == BackendType.TORCH:
+        from torch import nn
 
-            if isinstance(model, GraphModelWrapper):
-                return model.get_graph()
-            msg = f"Unexpected type of model {type(model)} for TORCH backend"
-            raise nncf.InternalError(msg)
-        msg = f"Cannot create backend-specific graph because {model_backend.value} is not supported!"
-        raise nncf.UnsupportedBackendError(msg)
+        from nncf.torch.function_hook.nncf_graph.nncf_graph_builder import GraphModelWrapper
+        from nncf.torch.function_hook.nncf_graph.nncf_graph_builder import build_nncf_graph
+
+        if isinstance(model, nn.Module):
+            return build_nncf_graph(model, example_input)
+
+        if isinstance(model, GraphModelWrapper):
+            return model.get_graph()
+
+        msg = f"Unexpected type of model {type(model)} for TORCH backend"
+        raise nncf.InternalError(msg)
+
+    msg = f"Cannot create backend-specific graph because {model_backend.value} is not supported!"
+    raise nncf.UnsupportedBackendError(msg)
 
 
 class ModelTransformerFactory:
@@ -165,7 +177,7 @@ def create(model: TModel) -> CommandCreator:
 
 class StatisticsAggregatorFactory:
     @staticmethod
-    def create(model: TModel, dataset: Dataset) -> aggregator.StatisticsAggregator:
+    def create(model: TModel, dataset: Dataset) -> StatisticsAggregator:
         """
         Factory method to create backend-specific `StatisticsAggregator` instance based on the input model.
 

src/nncf/experimental/torch/fx/quantization/quantize_model.py

@@ -21,7 +21,7 @@
 from torch.fx.passes.infra.pass_manager import PassManager
 
 import nncf
-from nncf.common.factory import NNCFGraphFactory
+from nncf.common.factory import build_graph
 from nncf.common.logging import nncf_logger
 from nncf.common.quantization.structs import QuantizationPreset
 from nncf.data import Dataset
@@ -90,7 +90,7 @@ def quantize_impl(
     # To make it easier for bias correction algorithms.
     apply_quantization_transformations(copied_model)
 
-    nncf_graph = NNCFGraphFactory.create(copied_model)
+    nncf_graph = build_graph(copied_model)
     quantized_model = quantization_algorithm.apply(copied_model, nncf_graph, dataset=calibration_dataset)
 
     if is_weight_compression_needed(advanced_parameters):
@@ -154,7 +154,7 @@ def compress_weights_impl(
         compression_format,
         advanced_parameters,
     )
-    graph = NNCFGraphFactory.create(model)
+    graph = build_graph(model)
     compressed_model = compression_algorithm.apply(model, graph, dataset=dataset)
     compressed_model = GraphModule(compressed_model, compressed_model.graph)
     compressed_model = _disallow_eval_train(compressed_model)

src/nncf/experimental/torch/fx/quantization/quantize_pt2e.py

@@ -25,7 +25,7 @@
 from nncf import AdvancedCompressionParameters
 from nncf import Dataset
 from nncf import SensitivityMetric
-from nncf.common.factory import NNCFGraphFactory
+from nncf.common.factory import build_graph
 from nncf.common.logging import nncf_logger
 from nncf.common.utils.api_marker import api
 from nncf.experimental.quantization.algorithms.post_training.algorithm import ExperimentalPostTrainingQuantization
@@ -123,7 +123,7 @@ def quantize_pt2e(
         batchwise_statistics=batchwise_statistics,
     )
 
-    nncf_graph = NNCFGraphFactory.create(transformed_model)
+    nncf_graph = build_graph(transformed_model)
     quantized_model = quantization_algorithm.apply(transformed_model, nncf_graph, dataset=calibration_dataset)
 
     # Magic. Without this call compiled model is not performant
@@ -225,7 +225,7 @@ def compress_pt2e(
 
     # Here the model is annotated
     transformed_model = quantizer.transform_prior_quantization(model)
-    nncf_graph = NNCFGraphFactory.create(transformed_model)
+    nncf_graph = build_graph(transformed_model)
     quantized_model = quantization_algorithm.apply(transformed_model, nncf_graph, dataset=dataset)
     quantized_model = torch.fx.GraphModule(quantized_model, graph=quantized_model.graph)
     return quantized_model

src/nncf/experimental/torch/sparsify_activations/sparsify_activations_impl.py

@@ -15,7 +15,7 @@
 
 import nncf
 from nncf.common import factory
-from nncf.common.factory import NNCFGraphFactory
+from nncf.common.factory import build_graph
 from nncf.common.graph.graph import NNCFGraph
 from nncf.common.graph.graph import NNCFNode
 from nncf.common.graph.operator_metatypes import OperatorMetatype
@@ -258,7 +258,7 @@ def sparsify_activations(
 
     algorithm = SparsifyActivationsAlgorithm(target_sparsity_by_scope, ignored_scope)
 
-    graph = NNCFGraphFactory.create(model)
+    graph = build_graph(model)
     sparse_model = algorithm.apply(model, graph, dataset)
 
     if backend == BackendType.TORCH:

src/nncf/onnx/graph/model_transformer.py

@@ -16,7 +16,7 @@
 import onnx
 
 import nncf
-from nncf.common.factory import NNCFGraphFactory
+from nncf.common.factory import build_graph
 from nncf.common.graph.model_transformer import ModelTransformer
 from nncf.common.graph.transformations.commands import TargetType
 from nncf.common.graph.transformations.layout import TransformationLayout
@@ -503,7 +503,7 @@ def _apply_multiply_insertion_transformations(
         """
         node_name_to_node = get_name_to_node_map(model)
         # TODO(andrey-churkin): Optimize it
-        graph = NNCFGraphFactory.create(model)
+        graph = build_graph(model)
         input_edges_mapping = get_input_edges_mapping(graph)
 
         for transformation in transformations:

src/nncf/onnx/quantization/quantize_model.py

@@ -21,7 +21,7 @@
 from onnx.external_data_helper import uses_external_data
 
 import nncf
-from nncf.common.factory import NNCFGraphFactory
+from nncf.common.factory import build_graph
 from nncf.common.logging.logger import nncf_logger
 from nncf.common.quantization.structs import QuantizationPreset
 from nncf.data import Dataset
@@ -360,7 +360,7 @@ def compress_weights_impl(
         compression_format,
         advanced_parameters,
     )
-    graph = NNCFGraphFactory.create(model)
+    graph = build_graph(model)
 
     compressed_model = compression_algorithm.apply(model, graph, dataset=dataset)
 

src/nncf/openvino/quantization/quantize_model.py

@@ -16,8 +16,8 @@
 import openvino as ov
 from openvino._offline_transformations import compress_quantize_weights_transformation
 
-from nncf.common.factory import NNCFGraphFactory
 from nncf.common.factory import StatisticsAggregatorFactory
+from nncf.common.factory import build_graph
 from nncf.common.logging import nncf_logger
 from nncf.common.quantization.structs import QuantizationPreset
 from nncf.data import Dataset
@@ -87,7 +87,7 @@ def _extract_all_subgraphs(model: ov.Model, current_id: str) -> None:
         :param model: Model.
         :param current_id: Current graph id.
         """
-        graphs[current_id] = NNCFGraphFactory.create(model)
+        graphs[current_id] = build_graph(model)
         for op in model.get_ops():
             if get_node_metatype(op) == OVIfMetatype:
                 _extract_all_subgraphs(op.get_function(0), op.get_friendly_name() + "_then")
@@ -384,7 +384,7 @@ def compress_weights_impl(
     Implementation of the `compress_weights()` method for the OpenVINO backend.
     """
     model = remove_friendly_name_duplicates(model)
-    graph = NNCFGraphFactory.create(model)
+    graph = build_graph(model)
     compression_algorithm = WeightCompression(
         mode,
         ratio,

src/nncf/quantization/algorithms/accuracy_control/algorithm.py

@@ -13,7 +13,7 @@
 from typing import Iterable, Optional, TypeVar
 
 import nncf
-from nncf.common.factory import NNCFGraphFactory
+from nncf.common.factory import build_graph
 from nncf.common.graph import NNCFGraph
 from nncf.common.graph import NNCFNode
 from nncf.common.graph.utils import get_number_of_quantized_ops
@@ -252,8 +252,8 @@ def _apply(
         :return: The quantized model whose metric `final_metric` is satisfied
             the maximum accuracy drop condition.
         """
-        initial_model_graph = NNCFGraphFactory.create(initial_model)
-        quantized_model_graph = NNCFGraphFactory.create(quantized_model)
+        initial_model_graph = build_graph(initial_model)
+        quantized_model_graph = build_graph(quantized_model)
 
         # Collect original biases and weights because these values are
         # required to undo bias correction and weight correction.

src/nncf/quantization/algorithms/bias_correction/algorithm.py

@@ -17,7 +17,7 @@
 from nncf import nncf_logger
 from nncf.common.factory import EngineFactory
 from nncf.common.factory import ModelTransformerFactory
-from nncf.common.factory import NNCFGraphFactory
+from nncf.common.factory import build_graph
 from nncf.common.graph import NNCFGraph
 from nncf.common.graph import NNCFNode
 from nncf.common.graph.definitions import NNCFGraphNodeType
@@ -143,9 +143,9 @@ def apply(
         main_model_transformer = ModelTransformerFactory.create(model)
 
         model_copy = copy_model(model)
-        graph_copy = NNCFGraphFactory.create(model_copy)
+        graph_copy = build_graph(model_copy)
         model_copy = self._backend_entity.remove_fq_from_inputs(model_copy, graph_copy)
-        nncf_graph = NNCFGraphFactory.create(model_copy)
+        nncf_graph = build_graph(model_copy)
 
         nodes_with_bias = []
         for node in nncf_graph.topological_sort():

src/nncf/quantization/algorithms/hyperparameter_tuner/algorithm.py

@@ -16,7 +16,7 @@
 import operator
 from typing import Any, Callable, Iterable, Mapping, Optional, TypeVar, Union
 
-from nncf.common.factory import NNCFGraphFactory
+from nncf.common.factory import build_graph
 from nncf.common.graph.graph import NNCFGraph
 from nncf.common.logging import nncf_logger
 from nncf.common.tensor_statistics.statistic_point import StatisticPointsContainer
@@ -284,7 +284,7 @@ def apply(self, model: TModel, validation_dataset: Dataset) -> TModel:
         best_settings = {}
 
         for step_index, step_param_grid in enumerate(self._param_grids):
-            step_graph = NNCFGraphFactory.create(step_model)
+            step_graph = build_graph(step_model)
 
             # If there are no parameters to optimize for the current step, simply execute
             # this pipeline step on the model.

src/nncf/quantization/algorithms/pipeline.py

@@ -11,8 +11,8 @@
 
 from typing import Optional, TypeVar, Union
 
-from nncf.common.factory import NNCFGraphFactory
 from nncf.common.factory import StatisticsAggregatorFactory
+from nncf.common.factory import build_graph
 from nncf.common.graph.graph import NNCFGraph
 from nncf.common.logging import nncf_logger
 from nncf.common.tensor_statistics.statistic_point import StatisticPointsContainer
@@ -115,7 +115,7 @@ def run_step(
         pipeline_step = pipeline_steps[step_index]
         for algorithm in pipeline_step[:-1]:
             current_model = algorithm.apply(current_model, current_graph, step_statistics)
-            current_graph = NNCFGraphFactory.create(current_model)
+            current_graph = build_graph(current_model)
         current_model = pipeline_step[-1].apply(current_model, current_graph, step_statistics)
 
         return current_model
@@ -152,7 +152,7 @@ def run_from_step(
         for step_index in range(start_step_index, len(pipeline_steps)):
             # Create graph required to run current pipeline step
             if step_graph is None:
-                step_graph = NNCFGraphFactory.create(step_model)
+                step_graph = build_graph(step_model)
 
             # Collect statistics required to run current pipeline step
             step_statistics = step_index_to_statistics.get(step_index)

src/nncf/torch/quantization/quantize_model.py

@@ -15,7 +15,7 @@
 import torch
 
 import nncf
-from nncf.common.factory import NNCFGraphFactory
+from nncf.common.factory import build_graph
 from nncf.common.quantization.structs import QuantizationPreset
 from nncf.data import Dataset
 from nncf.parameters import BackupMode
@@ -123,7 +123,7 @@ def compress_weights_impl(
         compression_format,
         advanced_parameters,
     )
-    graph = NNCFGraphFactory.create(model)
+    graph = build_graph(model)
 
     compressed_model = compression_algorithm.apply(model, graph, dataset=dataset)
     if isinstance(compressed_model, GraphModelWrapper):