Test ResNet Ops Support

.github/workflows/precommit.yml

@@ -208,7 +208,7 @@ jobs:
           lfs: true
       - uses: actions/setup-python@0a5c61591373683505ea898e09a3ea4f39ef2b9c # v5.0.0
         with:
-          python-version: 3.8.18
+          python-version: 3.9.19
           cache: pip
       - name: Install NNCF and test requirements
         run: make install-tensorflow-test

README.md

@@ -429,7 +429,7 @@ conda install -c conda-forge nncf
 - Python\* 3.8 or later
 - Supported frameworks:
   - PyTorch\* >=2.2, <2.4
-  - TensorFlow\* >=2.8.4, <=2.12.1
+  - TensorFlow\* >=2.8.4, <=2.15.1
   - ONNX\* ==1.16.0
   - OpenVINO\* >=2022.3.0
 

constraints.txt

@@ -10,7 +10,8 @@ onnx==1.16.0
 onnxruntime==1.17.1
 
 # TensorFlow
-tensorflow==2.12.1
+tensorflow==2.12.1; python_version < '3.9'
+tensorflow==2.15.1; python_version >= '3.9'
 
 # Tests and examples
 pytest==8.0.2

docs/Installation.md

@@ -43,7 +43,7 @@ as well as the supported versions of Python:
 
 | NNCF      | OpenVINO   | PyTorch  | ONNX     | TensorFlow | Python |
 |-----------|------------|----------|----------|------------|--------|
-| `develop` | `2024.2.0` | `2.3.0`  | `1.16.0` | `2.12.0`   | `3.8`  |
+| `develop` | `2024.2.0` | `2.3.0`  | `1.16.0` | `2.15.1`   | `3.8`* |
 | `2.11.0`  | `2024.2.0` | `2.3.0`  | `1.16.0` | `2.12.0`   | `3.8`  |
 | `2.10.0`  | `2024.1.0` | `2.2.1`  | `1.16.0` | `2.12.0`   | `3.8`  |
 | `2.9.0`   | `2024.0.0` | `2.1.2`  | `1.13.1` | `2.12.0`   | `3.8`  |
@@ -53,3 +53,5 @@ as well as the supported versions of Python:
 | `2.6.0`   | `2023.1.0` | `2.0.1`  | `1.13.1` | `2.12.0`   | `3.8`  |
 | `2.5.0`   | `2023.0.0` | `1.13.1` | `1.13.1` | `2.11.1`   | `3.8`  |
 | `2.4.0`   | `2022.1.0` | `1.12.1` | `1.12.0` | `2.8.2`    | `3.8`  |
+
+> (*) Python 3.9 or higher is required for TensorFlow 2.15.1

examples/post_training_quantization/tensorflow/mobilenet_v2/requirements.txt

@@ -1,4 +1,5 @@
-tensorflow~=2.12.0
+tensorflow~=2.12.0; python_version < '3.9'
+tensorflow~=2.15.1; python_version >= '3.9'
 tensorflow-datasets
 tqdm
 openvino==2024.2

examples/tensorflow/common/object_detection/architecture/darknet.py

@@ -11,9 +11,12 @@
 
 import tensorflow as tf
 import tensorflow.keras.backend as K
+from packaging import version
 
 from examples.tensorflow.common.object_detection.architecture import nn_ops
 
+tensorflow_version = version.parse(version.parse(tf.__version__).base_version)
+
 
 class CSPDarknet53:
     """Class to build CSPDarknet53"""
@@ -25,12 +28,17 @@ def DarknetConv2D_BN_Mish(self, *args, **kwargs):
         """Darknet Convolution2D followed by SyncBatchNormalization and Mish."""
         no_bias_kwargs = {"use_bias": False}
         no_bias_kwargs.update(kwargs)
+
+        if tensorflow_version < version.parse("2.15"):
+            mish = tf.keras.layers.Activation(self.mish)
+        else:
+            mish = tf.keras.layers.Activation("mish")
+
         return nn_ops.compose(
             nn_ops.DarknetConv2D(*args, **no_bias_kwargs),
             # TODO(nsavelyev) replace by BatchNormalization(synchronized=True) once support for TF < 2.12 is dropped
             tf.keras.layers.experimental.SyncBatchNormalization(),
-            # TODO(nsavelyev) change to tf.keras.activations.mish after upgrade to TF 2.13
-            tf.keras.layers.Activation(self.mish),
+            mish,
         )
 
     def csp_resblock_body(self, x, num_filters, num_blocks, all_narrow=True):

examples/tensorflow/requirements.txt

@@ -4,7 +4,8 @@ absl-py==1.0.0
 tensorflow
 tensorflow_datasets==4.2.0
 tensorflow_hub
-tensorflow_addons==0.20.0
+tensorflow_addons==0.20.0; python_version < '3.9'
+tensorflow_addons==0.23.0; python_version >= '3.9'
 tensorflow-metadata==1.13.0
 opencv-python
 pycocotools==2.0.6

nncf/common/factory.py

@@ -41,6 +41,10 @@ def create(model: TModel) -> NNCFGraph:
         if model_backend == BackendType.OPENVINO:
             from nncf.openvino.graph.nncf_graph_builder import GraphConverter
 
+            return GraphConverter.create_nncf_graph(model)
+        if model_backend == BackendType.TORCH_FX:
+            from nncf.experimental.torch.fx.nncf_graph_builder import GraphConverter
+
             return GraphConverter.create_nncf_graph(model)
         if model_backend == BackendType.TORCH:
             return model.nncf.get_graph()
@@ -72,6 +76,10 @@ def create(model: TModel, inplace: bool = False) -> ModelTransformer:
             from nncf.torch.model_transformer import PTModelTransformer
 
             return PTModelTransformer(model)
+        if model_backend == BackendType.TORCH_FX:
+            from nncf.experimental.torch.fx.model_transformer import FXModelTransformer
+
+            return FXModelTransformer(model)
         raise nncf.UnsupportedBackendError(
             "Cannot create backend-specific model transformer because {} is not supported!".format(model_backend.value)
         )
@@ -95,7 +103,7 @@ def create(model: TModel) -> Engine:
             from nncf.openvino.engine import OVNativeEngine
 
             return OVNativeEngine(model)
-        if model_backend == BackendType.TORCH:
+        if model_backend in (BackendType.TORCH, BackendType.TORCH_FX):
             from nncf.torch.engine import PTEngine
 
             return PTEngine(model)
@@ -151,6 +159,10 @@ def create(model: TModel, dataset: Dataset) -> aggregator.StatisticsAggregator:
             from nncf.torch.statistics.aggregator import PTStatisticsAggregator
 
             return PTStatisticsAggregator(dataset)
+        if model_backend == BackendType.TORCH_FX:
+            from nncf.experimental.torch.fx.statistics.aggregator import FXStatisticsAggregator
+
+            return FXStatisticsAggregator(dataset)
         raise nncf.UnsupportedBackendError(
             "Cannot create backend-specific statistics aggregator because {} is not supported!".format(
                 model_backend.value

nncf/common/graph/patterns/manager.py

@@ -45,7 +45,7 @@ def _get_backend_hw_patterns_map(backend: BackendType) -> Dict[HWFusedPatternNam
 
             registry = OPENVINO_HW_FUSED_PATTERNS.registry_dict
             return registry
-        if backend == BackendType.TORCH:
+        if backend in (BackendType.TORCH, BackendType.TORCH_FX):
             from nncf.torch.hardware.fused_patterns import PT_HW_FUSED_PATTERNS
 
             registry = PT_HW_FUSED_PATTERNS.registry_dict
@@ -73,7 +73,7 @@ def _get_backend_ignored_patterns_map(
 
             registry = OPENVINO_IGNORED_PATTERNS.registry_dict
             return registry
-        if backend == BackendType.TORCH:
+        if backend in (BackendType.TORCH, BackendType.TORCH_FX):
             from nncf.torch.quantization.ignored_patterns import PT_IGNORED_PATTERNS
 
             registry = PT_IGNORED_PATTERNS.registry_dict

nncf/common/utils/backend.py

@@ -20,6 +20,7 @@
 
 class BackendType(Enum):
     TORCH = "Torch"
+    TORCH_FX = "TorchFX"
     TENSORFLOW = "Tensorflow"
     ONNX = "ONNX"
     OPENVINO = "OpenVINO"
@@ -33,6 +34,7 @@ def get_available_backends() -> List[BackendType]:
     """
     frameworks = [
         ("torch", BackendType.TORCH),
+        ("torch.fx", BackendType.TORCH_FX),
         ("tensorflow", BackendType.TENSORFLOW),
         ("onnx", BackendType.ONNX),
         ("openvino.runtime", BackendType.OPENVINO),
@@ -51,14 +53,27 @@ def get_available_backends() -> List[BackendType]:
 
 def is_torch_model(model: TModel) -> bool:
     """
-    Returns True if the model is an instance of torch.nn.Module, otherwise False.
+    Returns True if the model is an instance of torch.nn.Module and not a torch.fx.GraphModule, otherwise False.
 
     :param model: A target model.
-    :return: True if the model is an instance of torch.nn.Module, otherwise False.
+    :return: True if the model is an instance of torch.nn.Module and not torch.fx.GraphModule, otherwise False.
     """
     import torch
+    import torch.fx
 
-    return isinstance(model, torch.nn.Module)
+    return not isinstance(model, torch.fx.GraphModule) and isinstance(model, torch.nn.Module)
+
+
+def is_torch_fx_model(model: TModel) -> bool:
+    """
+    Returns True if the model is an instance of torch.fx.GraphModule, otherwise False.
+
+    :param model: A target model.
+    :return: True if the model is an instance of torch.fx.GraphModule, otherwise False.
+    """
+    import torch.fx
+
+    return isinstance(model, torch.fx.GraphModule)
 
 
 def is_tensorflow_model(model: TModel) -> bool:
@@ -118,6 +133,9 @@ def get_backend(model: TModel) -> BackendType:
     """
     available_backends = get_available_backends()
 
+    if BackendType.TORCH_FX in available_backends and is_torch_fx_model(model):
+        return BackendType.TORCH_FX
+
     if BackendType.TORCH in available_backends and is_torch_model(model):
         return BackendType.TORCH
 

nncf/experimental/torch/fx/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) 2024 Intel Corporation
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#      http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.

nncf/experimental/torch/fx/model_transformer.py

@@ -0,0 +1,112 @@
+# Copyright (c) 2024 Intel Corporation
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#      http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 collections import defaultdict
+from typing import Callable, List, Union
+
+import torch
+import torch.fx
+from torch.fx.passes.split_utils import split_by_tags
+
+from nncf.common.graph.model_transformer import ModelTransformer
+from nncf.common.graph.transformations.commands import Command
+from nncf.common.graph.transformations.commands import TransformationPriority
+from nncf.common.graph.transformations.commands import TransformationType
+from nncf.torch.graph.transformations.commands import PTModelExtractionCommand
+from nncf.torch.graph.transformations.layout import PTTransformationLayout
+
+
+class FXApplyTransformationCommand(Command):
+    def __init__(
+        self,
+        transformation_fn: Callable[[torch.fx.GraphModule], None],
+        priority: Union[TransformationPriority, int] = TransformationPriority.DEFAULT_PRIORITY,
+    ):
+        super().__init__(TransformationType.INSERT)
+        self.tranformation_fn = transformation_fn
+        self.priority = priority
+
+
+class FXModelTransformer(ModelTransformer):
+    """
+    Applies transformations upon Torch FX model.
+    """
+
+    def __init__(self, model: torch.fx.GraphModule):
+        super().__init__(model)
+
+        self._command_transformation_ordered_pairs = [
+            (FXApplyTransformationCommand, self._apply_transformation),
+            (PTModelExtractionCommand, self._apply_model_extraction),
+        ]
+
+    def transform(self, transformation_layout: PTTransformationLayout) -> torch.fx.GraphModule:
+        # TODO(dlyakhov): Manage priorities of transformations.
+        transformations = transformation_layout.transformations
+        aggregated_transformations = defaultdict(list)
+        for transformation in transformations:
+            aggregated_transformations[transformation.__class__].append(transformation)
+
+        model = self._model
+        for transformation_cls, transformation_fn in self._command_transformation_ordered_pairs:
+            transformations = aggregated_transformations[transformation_cls]
+            if transformations:
+                model = transformation_fn(model, transformations)
+
+        # Do not use model.graph.eliminate_dead_code()
+        # because the computational statistics code
+        # is interpolated as dead code.
+        model.recompile()
+        return model
+
+    @staticmethod
+    def _apply_model_extraction(
+        model: torch.fx.GraphModule,
+        transformations: List[PTModelExtractionCommand],
+    ) -> torch.fx.GraphModule:
+        transformation = transformations[-1]
+        assert len(transformation.input_node_names) == 1
+        assert transformation.input_node_names == transformation.output_node_names
+        node_name = transformation.input_node_names[0]
+
+        tags = ["before", "extracted", "after"]
+        i = 0
+        for node in model.graph.nodes:
+            if node.name == node_name:
+                node.tag = tags[1]
+                weights = [node.all_input_nodes[1]]
+                while weights:
+                    w_node = weights.pop()
+                    assert w_node.tag in tags[0:2]
+                    w_node.tag = tags[1]
+                    weights.extend(w_node.all_input_nodes)
+                i = 2
+                continue
+            node.tag = tags[i]
+
+        splitted_gm = split_by_tags(model, tags)
+        return splitted_gm.extracted
+
+    @staticmethod
+    def get_graph_node_by_name(graph: torch.fx.Graph, name: str) -> torch.fx.Node:
+        for node in graph.nodes:
+            if node.name == name:
+                return node
+        raise RuntimeError(f"Node with name {name} is not found")
+
+    @staticmethod
+    def _apply_transformation(
+        model: torch.fx.GraphModule,
+        transformations: List[FXApplyTransformationCommand],
+    ) -> torch.fx.GraphModule:
+        for transformation in transformations:
+            transformation.tranformation_fn(model)
+        return model