[TF] fix custom layers sparsification, add graph test (openvinotoolki…

…t#522)

* fix custom layers sparsification

beta/README.md

@@ -91,7 +91,7 @@ Quick jump to the samples:
 |ResNet50|Sparsity 50% (Magnitude)|ImageNet|75.04|75|
 |ResNet50|INT8 w:sym,per-tensor a:sym,per-tensor + Sparsity 50% (Magnitude)|ImageNet|75.04|74.46|
 |ResNet50|Filter Pruning 40%|ImageNet|75.04|74.98|
-|TensorFlow Hub MobileNet V2|Sparsity 35% (Magnitude)|ImageNet|71.84|71.73|
+|TensorFlow Hub MobileNet V2|Sparsity 35% (Magnitude)|ImageNet|71.84|71.90|
 
 #### Object detection
 

beta/examples/tensorflow/classification/README.md

@@ -112,4 +112,4 @@ To export a model to the OpenVINO IR and run it using the Intel® Deep Learning
 |ResNet50|Sparsity 50% (Magnitude)|ImageNet|75|[resnet50_imagenet_magnitude_sparsity.json](configs/sparsity/resnet50_imagenet_magnitude_sparsity.json)|[Link](https://storage.openvinotoolkit.org/repositories/nncf/tensorflow/models/develop/resnet50_sparsity_50.tar.gz)|
 |ResNet50|INT8 w:sym,per-tensor a:sym,per-tensor + Sparsity 50% (Magnitude)|ImageNet|74.46|[resnet50_imagenet_magnitude_sparsity_int8.json](configs/sparsity_quantization/resnet50_imagenet_magnitude_sparsity_int8.json)|[Link](https://storage.openvinotoolkit.org/repositories/nncf/tensorflow/models/develop/resnet50_int8_w_sym_t_a_sym_t_sparsity_50.tar.gz)|
 |ResNet50|Filter Pruning 40%|ImageNet|74.98|[resnet50_imagenet_pruning_geometric_median.json](configs/pruning/resnet50_imagenet_pruning_geometric_median.json)|[Link](https://storage.openvinotoolkit.org/repositories/nncf/tensorflow/models/develop/resnet50_pruning_40.tar.gz)|
-|TensorFlow Hub MobileNet V2|Sparsity 35% (Magnitude)|ImageNet|71.83|[mobilenet_v2_hub_imagenet_magnitude_sparsity.json](configs/sparsity/mobilenet_v2_hub_imagenet_magnitude_sparsity.json)|[Link](https://storage.openvinotoolkit.org/repositories/nncf/tensorflow/models/develop/tf1_mobilenet_v2_1.0_224_s0.35.tar.gz)|
+|TensorFlow Hub MobileNet V2|Sparsity 35% (Magnitude)|ImageNet|71.90|[mobilenet_v2_hub_imagenet_magnitude_sparsity.json](configs/sparsity/mobilenet_v2_hub_imagenet_magnitude_sparsity.json)|[Link](https://storage.openvinotoolkit.org/repositories/nncf/tensorflow/models/develop/tf1_mobilenet_v2_1.0_224_s0.35.tar.gz)|

beta/examples/tensorflow/classification/configs/sparsity/mobilenet_v2_hub_imagenet_magnitude_sparsity.json

@@ -12,8 +12,8 @@
         "type": "Adam",
         "schedule_type": "piecewise_constant",
         "schedule_params": {
-            "boundaries": [12, 14, 16],
-            "values":  [1e-5, 1e-6, 1e-7, 1e-8]
+            "boundaries": [14],
+            "values":  [1e-5, 1e-6]
         }
     },
 

beta/nncf/tensorflow/sparsity/magnitude/algorithm.py

@@ -50,6 +50,13 @@
     'LocallyConnected2D': {WEIGHT_ATTR_NAME: 'kernel'}
 }
 
+SPARSITY_TF_OPS = [
+    'Conv2D',
+    'Conv3D',
+    'DepthwiseConv2dNative',
+    'QuantizedConv2D'
+]
+
 
 @TF_COMPRESSION_ALGORITHMS.register('magnitude_sparsity')
 class MagnitudeSparsityBuilder(TFCompressionAlgorithmBuilder):
@@ -83,7 +90,7 @@ def get_transformation_layout(self, model):
         for layer in get_custom_layers(model):
             nxmodel = convert_layer_graph_to_nxmodel(layer)
             for node_name, node in nxmodel.nodes.items():
-                if node['type'] in SPARSITY_LAYERS \
+                if node['type'] in SPARSITY_TF_OPS \
                         and not is_ignored(node_name, self.ignored_scopes):
                     weight_attr_name = get_weight_node_name(nxmodel, node_name)
                     transformations.register(

beta/tests/sota_checkpoints_eval.json

@@ -87,7 +87,7 @@
                 },
                 "tf1_mobilenet_v2_1.0_224_sparsity_35": {
                     "config": "examples/tensorflow/classification/configs/sparsity/mobilenet_v2_hub_imagenet_magnitude_sparsity.json",
-                    "target": 71.83,
+                    "target": 71.90,
                     "resume": "tf1_mobilenet_v2_1.0_224_sparsity_35",
                     "metric_type": "Acc@1",
                     "model_description": "tf1 mobilenet_v2_sparsity_35",

beta/tests/tensorflow/test_compressed_graph.py

@@ -38,21 +38,22 @@ def get_basic_quantization_config(qconfig, input_sample_sizes=None):
     return config
 
 
-def get_nx_graph_from_tf_graph(tf_graph: tf.Graph):
+def get_nx_graph_from_tf_graph(tf_graph: tf.Graph, graph_to_layer_var_names_map: dict):
     def _get_node_attributes(op: tf.Operation):
         attr = {'op': op.type}
         return attr
 
     def _get_inbound_edges(op: tf.Operation):
         inbound_edges = []
         for input_tensor in op.inputs:
-            inbound_edges.append((input_tensor.op.name, op.name))
+            inbound_edges.append((graph_to_layer_var_names_map.get(input_tensor.op.name, input_tensor.op.name),
+                                  graph_to_layer_var_names_map.get(op.name, op.name)))
         return inbound_edges
 
     nodes = {}
     edges = []
     for op in tf_graph.get_operations():
-        op_name = op.name
+        op_name = graph_to_layer_var_names_map.get(op.name, op.name)
         nodes[op_name] = _get_node_attributes(op)
         edges.extend(_get_inbound_edges(op))
 
@@ -87,38 +88,6 @@ def check_nx_graph(nx_graph: nx.DiGraph, graph_path: str):
     assert nx.DiGraph(expected_graph).edges == nx_graph.edges
 
 
-def check_graph(tf_graph: tf.Graph, ref_graph_dir: str, ref_graph_filename: str):
-    data_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data', 'reference_graphs')
-    graph_dir = os.path.join(data_dir, ref_graph_dir)
-    graph_path = os.path.abspath(os.path.join(graph_dir, ref_graph_filename))
-
-    # validate file with graph manually!
-    ref_graph_not_exist = False
-    if not os.path.exists(graph_path):
-        if not os.path.exists(graph_dir):
-            os.makedirs(graph_dir)
-        ref_graph_not_exist = True
-
-    _, ref_graph_ext = os.path.splitext(ref_graph_filename)
-
-    if ref_graph_ext == '.pb':
-        graph_def = tf_graph.as_graph_def(add_shapes=True)
-        # remove control edges for a human-readable graph visualization
-        # remove_control_edges(graph_def)
-
-        if ref_graph_not_exist:
-            tf.io.write_graph(graph_def, graph_dir, ref_graph_filename, as_text=False)
-
-        check_graph_def(graph_def, graph_path)
-    else:
-        nx_graph = get_nx_graph_from_tf_graph(tf_graph)
-
-        if ref_graph_not_exist:
-            nx.drawing.nx_pydot.write_dot(nx_graph, graph_path)
-
-        check_nx_graph(nx_graph, graph_path)
-
-
 class QuantizeTestCaseConfiguration:
     def __init__(self, quant_mode, quant_granularity, graph_dir):
         self.qconfig = Dict()
@@ -178,17 +147,20 @@ def _pruning_case_config(request):
 
 
 class ModelDesc:
-    def __init__(self, ref_graph_filename: str, model_builder, input_sample_sizes):
+    def __init__(self, ref_graph_filename: str, model_builder, input_sample_sizes,
+                 rename_resource_nodes=False):
         self.model_name, _ = os.path.splitext(ref_graph_filename)
         self.model_builder = model_builder
         self.ref_graph_filename = ref_graph_filename
         self.input_sample_sizes = input_sample_sizes
+        self.rename_resource_nodes = rename_resource_nodes
 
 
 SKIP_MAP = {
     'quantization': {
         'inception_resnet_v2': pytest.mark.skip(reason='gitlab issue #17'),
         'nasnet_mobile': pytest.mark.skip(reason='gitlab issue #18'),
+        'mobilenet_v2_slim': pytest.mark.skip(reason='ticket #46349'),
         'xception': pytest.mark.skip(reason='gitlab issue #28')
     },
     'magnitude_sparsity': {
@@ -204,6 +176,7 @@ def __init__(self, ref_graph_filename: str, model_builder, input_sample_sizes):
         'mask_rcnn': pytest.mark.skip(reason='ticket #50605'),
         'mobilenet_v3_small': pytest.mark.skip(reason='ticket #50607'),
         'yolo_v4': pytest.mark.skip(reason='ticket #50608'),
+        'mobilenet_v2_slim': pytest.mark.skip(reason='ticket #46349'),
     }
 }
 
@@ -254,6 +227,10 @@ def get_test_models_desc(algorithm):
             ModelDesc('yolo_v4.pb', test_models.YOLOv4, [1, None, None, 3]),
             marks=SKIP_MAP[algorithm].get('yolo_v4', ())
         ),
+        pytest.param(
+            ModelDesc('mobilenet_v2_slim.dot', test_models.HubMobileNetV2, [1, 224, 224, 3], True),
+            marks=SKIP_MAP[algorithm].get('mobilenet_v2_slim', ())
+        )
     ]
 
 
@@ -262,10 +239,29 @@ def keras_model_to_tf_graph(model):
     for item in model.inputs:
         input_signature.append(tf.TensorSpec(item.shape, item.dtype))
     concrete_function = tf.function(model).get_concrete_function(input_signature)
-    return concrete_function.graph
+    return concrete_function.graph, get_graph_to_layer_var_names_map(concrete_function)
+
+
+def get_graph_to_layer_var_names_map(concrete_fun):
+    names_map = {}
+    for layer_var in concrete_fun.variables:
+        for value_tensor, graph_name in concrete_fun.graph.captures:
+            if layer_var.handle is value_tensor:
+                names_map[graph_name.name.split(':')[0]] = layer_var.name.split(':')[0]
+    return names_map
+
+
+def rename_graph_def_nodes(graph_def, names_map: dict):
+    for node in graph_def.node:
+        node.name = names_map.get(node.name, node.name)
+        inp_names = []
+        for inp in node.input:
+            inp_names.append(names_map.get(inp, inp))
+        del node.input[:]
+        node.input.extend(inp_names)
 
 
-def remove_control_edges(graph_def):
+def remove_control_edges_from_graph_def(graph_def):
     for node in graph_def.node:
         inp_names = []
         for inp in node.input:
@@ -276,9 +272,59 @@ def remove_control_edges(graph_def):
         node.input.extend(inp_names)
 
 
-def check_model_graph(compressed_model, ref_graph_filename, ref_graph_dir):
-    compressed_graph = keras_model_to_tf_graph(compressed_model)
-    check_graph(compressed_graph, ref_graph_dir, ref_graph_filename)
+def prepare_and_check_graph_def(tf_graph: tf.Graph, graph_path: str,
+                                ref_graph_exist: bool,
+                                graph_to_layer_var_names_map=None,
+                                remove_control_edges=False):
+    graph_def = tf_graph.as_graph_def(add_shapes=True)
+    # remove control edges for a human-readable graph visualization
+    if remove_control_edges:
+        remove_control_edges_from_graph_def(graph_def)
+
+    if graph_to_layer_var_names_map:
+        rename_graph_def_nodes(graph_def, graph_to_layer_var_names_map)
+
+    if not ref_graph_exist:
+        graph_dir, ref_graph_filename = os.path.split(graph_path)
+        tf.io.write_graph(graph_def, graph_dir, ref_graph_filename, as_text=False)
+
+    check_graph_def(graph_def, graph_path)
+
+
+def prepare_and_check_nx_graph(tf_graph: tf.Graph, graph_path: str, ref_graph_exist: bool,
+                               graph_to_layer_var_names_map: dict):
+    nx_graph = get_nx_graph_from_tf_graph(tf_graph, graph_to_layer_var_names_map)
+
+    if not ref_graph_exist:
+        nx.drawing.nx_pydot.write_dot(nx_graph, graph_path)
+
+    check_nx_graph(nx_graph, graph_path)
+
+
+def check_model_graph(compressed_model, ref_graph_filename, ref_graph_dir, rename_resource_nodes):
+    data_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data', 'reference_graphs')
+    graph_dir = os.path.join(data_dir, ref_graph_dir)
+    graph_path = os.path.abspath(os.path.join(graph_dir, ref_graph_filename))
+
+    # validate file with graph manually!
+    ref_graph_exist = True
+    if not os.path.exists(graph_path):
+        if not os.path.exists(graph_dir):
+            os.makedirs(graph_dir)
+        ref_graph_exist = False
+
+    compressed_graph, graph_to_layer_var_names_map = keras_model_to_tf_graph(compressed_model)
+    if not rename_resource_nodes:
+        graph_to_layer_var_names_map = {}
+
+    ref_graph_ext = os.path.splitext(ref_graph_filename)[1]
+    if ref_graph_ext == '.pb':
+        prepare_and_check_graph_def(compressed_graph, graph_path, ref_graph_exist,
+                                    graph_to_layer_var_names_map)
+
+    else:
+        prepare_and_check_nx_graph(compressed_graph, graph_path, ref_graph_exist,
+                                   graph_to_layer_var_names_map)
 
 
 class TestModelsGraph:
@@ -294,7 +340,8 @@ def test_quantize_network(self, desc: ModelDesc, _quantization_case_config):
                                                input_sample_sizes=desc.input_sample_sizes)
         compressed_model, _ = create_compressed_model_and_algo_for_test(model, config)
 
-        check_model_graph(compressed_model, desc.ref_graph_filename, _quantization_case_config.graph_dir)
+        check_model_graph(compressed_model, desc.ref_graph_filename, _quantization_case_config.graph_dir,
+                          desc.rename_resource_nodes)
 
     @pytest.mark.parametrize(
         'desc', get_test_models_desc('magnitude_sparsity'), ids=[
@@ -308,7 +355,8 @@ def test_sparsity_network(self, desc: ModelDesc, _sparsity_case_config):
         config['compression']['params'] = {'schedule': 'multistep'}
         compressed_model, _ = create_compressed_model_and_algo_for_test(model, config)
 
-        check_model_graph(compressed_model, desc.ref_graph_filename, _sparsity_case_config.graph_dir)
+        check_model_graph(compressed_model, desc.ref_graph_filename, _sparsity_case_config.graph_dir,
+                          desc.rename_resource_nodes)
 
     @pytest.mark.parametrize(
         'desc', get_test_models_desc('filter_pruning'), ids=[
@@ -321,7 +369,8 @@ def test_pruning_network(self, desc: ModelDesc, _pruning_case_config):
         config = get_basic_filter_pruning_config(desc.input_sample_sizes)
         compressed_model, _ = create_compressed_model_and_algo_for_test(model, config)
 
-        check_model_graph(compressed_model, desc.ref_graph_filename, _pruning_case_config.graph_dir)
+        check_model_graph(compressed_model, desc.ref_graph_filename, _pruning_case_config.graph_dir,
+                          desc.rename_resource_nodes)
 
 
 QUANTIZE_OUTPUTS = [
@@ -340,4 +389,5 @@ def test_quantize_outputs(desc: ModelDesc, _quantization_case_config):
     config['compression']['quantize_outputs'] = True
     compressed_model, _ = create_compressed_model_and_algo_for_test(model, config)
 
-    check_model_graph(compressed_model, desc.ref_graph_filename, _quantization_case_config.graph_dir)
+    check_model_graph(compressed_model, desc.ref_graph_filename, _quantization_case_config.graph_dir,
+                      desc.rename_resource_nodes)

beta/tests/tensorflow/test_models/__init__.py

@@ -12,6 +12,7 @@
 """
 
 from beta.tests.tensorflow.test_models.densenet import DenseNet121
+from beta.tests.tensorflow.test_models.mobilenet_v2_slim import HubMobileNetV2
 from beta.tests.tensorflow.test_models.inception_resnet_v2 import InceptionResNetV2
 from beta.tests.tensorflow.test_models.inception_v3 import InceptionV3
 from beta.tests.tensorflow.test_models.mobilenet import MobileNet

beta/tests/tensorflow/test_models/mobilenet_v2_slim.py

@@ -0,0 +1,18 @@
+"""
+ Copyright (c) 2021 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 beta.examples.tensorflow.common.models import mobilenet_v2_100_224
+
+
+def HubMobileNetV2(input_shape=None):
+    return mobilenet_v2_100_224(input_shape)