Add tests for filter pruning by flops

nncf/tensorflow/pruning/filter_pruning/algorithm.py

@@ -353,7 +353,7 @@ def _set_binary_masks_for_pruned_modules_globally_by_flops_target(self,
         nncf_logger.debug('Setting new binary masks for pruned layers.')
         target_flops = self.full_flops * (1 - target_flops_pruning_rate)
         wrapped_layers = collect_wrapped_layers(self._model)
-        masks = []
+        masks = {}
 
         nncf_sorted_nodes = self._original_graph.topological_sort()
         for layer in wrapped_layers:
@@ -381,7 +381,7 @@ def _set_binary_masks_for_pruned_modules_globally_by_flops_target(self,
         for _, group_id, filter_index in sorted_importances:
             if self._pruning_quotas[group_id] == 0:
                 continue
-            masks[group_id][filter_index] = 0
+            masks[group_id] = tf.tensor_scatter_nd_update(masks[group_id], [[filter_index]], [0])
             self._pruning_quotas[group_id] -= 1
 
             # Update input/output shapes of pruned elements

tests/tensorflow/pruning/helpers.py

@@ -0,0 +1,106 @@
+"""
+ 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 addict import Dict
+
+from tensorflow.python.keras import layers
+import tensorflow as tf
+
+from nncf import NNCFConfig
+
+
+def get_basic_pruning_config(model_size=8):
+    config = NNCFConfig()
+    config.update(Dict({
+        "model": "basic",
+        "input_info":
+            {
+                "sample_size": [1, model_size, model_size, 1],
+            },
+        "compression":
+            {
+                "algorithm": "filter_pruning",
+                "pruning_init": 0.5,
+                "params": {
+                    "prune_first_conv": True,
+                    "prune_last_conv": True
+                }
+            }
+    }))
+    return config
+
+
+def get_concat_test_model(input_shape):
+    #             (input)
+    #                |
+    #             (conv1)
+    #        /       |       \
+    #    (conv2)  (conv3)  (conv4)
+    #       |        |       |
+    #       |    (gr_conv)   |
+    #         \    /         |
+    #        (concat)    (bn_conv4)
+    #             \       /
+    #              (concat)
+    #                 |
+    #            (bn_concat)
+    #                 |
+    #              (conv5)
+
+    inputs = tf.keras.Input(shape=input_shape[1:], name='input')
+    conv1 = layers.Conv2D(16, 1, name='conv1')
+    conv2 = layers.Conv2D(16, 1, name='conv2')
+    conv3 = layers.Conv2D(16, 1, name='conv3')
+    group_conv = layers.Conv2D(16, 1, groups=8, name='group_conv')
+    conv4 = layers.Conv2D(32, 1, name='conv4')
+    bn_conv4 = layers.BatchNormalization(name="bn_conv4")
+    bn_concat = layers.BatchNormalization(name="bn_concat")
+    conv5 = layers.Conv2D(64, 1, name='conv5')
+
+    x = conv1(inputs)
+    x1 = tf.concat([conv2(x), group_conv(conv3(x))], -1, name='tf_concat_1')
+    x = conv4(x)
+    x = bn_conv4(x)
+    x = tf.concat([x, x1], -1, name='tf_concat_2')
+    x = bn_concat(x)
+    outputs = conv5(x)
+    return tf.keras.Model(inputs=inputs, outputs=outputs)
+
+
+def init_conv_weights(weights):
+    new_weights = weights.numpy()
+    for i in range(weights.shape[-1]):
+        new_weights[..., i] += i
+    weights.assign(new_weights)
+
+
+def get_test_model_shared_convs(input_shape):
+    inputs = tf.keras.Input(shape=input_shape[1:], name='input')
+    conv1 = layers.Conv2D(512, 1, name='conv1', kernel_initializer='Ones',
+                          bias_initializer='Ones')
+    conv2 = layers.Conv2D(1024, 3, name='conv2', kernel_initializer='Ones',
+                          bias_initializer='Ones')
+    conv3 = layers.Conv2D(1024, 1, name='conv3', kernel_initializer='Ones',
+                          bias_initializer='Ones')
+    maxpool = layers.MaxPool2D()
+
+    in1 = conv1(inputs)
+    in2 = maxpool(in1)
+    out1 = conv2(in1)
+    out2 = conv2(in2)
+    x = conv3(out1)
+    y = conv3(out2)
+
+    init_conv_weights(conv1.kernel)
+    init_conv_weights(conv2.kernel)
+    init_conv_weights(conv3.kernel)
+    return tf.keras.Model(inputs=inputs, outputs=[x, y])

tests/tensorflow/pruning/test_algorithm.py

@@ -11,80 +11,19 @@
  limitations under the License.
 """
 
-from addict import Dict
-
-from tensorflow.python.keras import layers
-import tensorflow as tf
 import numpy as np
 import pytest
 
-from nncf import NNCFConfig
 from tests.tensorflow.helpers import create_compressed_model_and_algo_for_test
 from nncf.tensorflow.graph.utils import collect_wrapped_layers
-
-
-def get_basic_pruning_config(model_size=8):
-    config = NNCFConfig()
-    config.update(Dict({
-        "model": "basic",
-        "input_info":
-            {
-                "sample_size": [1, model_size, model_size, 1],
-            },
-        "compression":
-            {
-                "algorithm": "filter_pruning",
-                "pruning_init": 0.5,
-                "params": {
-                    "prune_first_conv": True,
-                    "prune_last_conv": True
-                }
-            }
-    }))
-    return config
+from tests.tensorflow.pruning.helpers import get_basic_pruning_config
+from tests.tensorflow.pruning.helpers import get_concat_test_model
 
 
 def check_pruning_mask(mask, pruning_rate, layer_name):
     assert np.sum(mask) == mask.size * pruning_rate, f"Incorrect masks for {layer_name}"
 
 
-def get_concat_test_model(input_shape):
-    #             (input)
-    #                |
-    #             (conv1)
-    #        /       |       \
-    #    (conv2)  (conv3)  (conv4)
-    #       |        |       |
-    #       |    (gr_conv)   |
-    #         \    /         |
-    #        (concat)    (bn_conv4)
-    #             \       /
-    #              (concat)
-    #                 |
-    #            (bn_concat)
-    #                 |
-    #              (conv5)
-
-    inputs = tf.keras.Input(shape=input_shape[1:], name='input')
-    conv1 = layers.Conv2D(16, 1, name='conv1')
-    conv2 = layers.Conv2D(16, 1, name='conv2')
-    conv3 = layers.Conv2D(16, 1, name='conv3')
-    group_conv = layers.Conv2D(16, 1, groups=8, name='group_conv')
-    conv4 = layers.Conv2D(32, 1, name='conv4')
-    bn_conv4 = layers.BatchNormalization(name="bn_conv4")
-    bn_concat = layers.BatchNormalization(name="bn_concat")
-    conv5 = layers.Conv2D(64, 1, name='conv5')
-
-    x = conv1(inputs)
-    x1 = tf.concat([conv2(x), group_conv(conv3(x))], -1, name='tf_concat_1')
-    x = conv4(x)
-    x = bn_conv4(x)
-    x = tf.concat([x, x1], -1, name='tf_concat_2')
-    x = bn_concat(x)
-    outputs = conv5(x)
-    return tf.keras.Model(inputs=inputs, outputs=outputs)
-
-
 @pytest.mark.parametrize(('all_weights', 'prune_batch_norms', 'ref_num_wrapped_layer'),
                          [
                              [True, True, 6],

tests/tensorflow/pruning/test_flops_pruning.py

@@ -0,0 +1,48 @@
+"""
+ 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.
+"""
+
+import pytest
+
+from tests.tensorflow.helpers import create_compressed_model_and_algo_for_test
+from tests.tensorflow.pruning.helpers import get_basic_pruning_config
+from tests.tensorflow.pruning.helpers import get_test_model_shared_convs
+
+
+@pytest.mark.parametrize(
+    ("model", "all_weights", "ref_full_flops", "ref_current_flops",
+     "ref_full_params", "ref_current_params"),
+    (
+        (get_test_model_shared_convs, True, 461438976, 276385312,
+         11534848, 6908711),
+        (get_test_model_shared_convs, False, 461438976, 270498816,
+         11534848, 6761608)
+    )
+)
+def test_flops_calulation_for_spec_layers(model, all_weights, ref_full_flops, ref_current_flops,
+                                          ref_full_params, ref_current_params):
+    config = get_basic_pruning_config(8)
+    config['compression']['algorithm'] = 'filter_pruning'
+    config['compression']['pruning_init'] = 0.4
+    config['compression']['params']['pruning_flops_target'] = 0.4
+    config['compression']['params']['prune_first_conv'] = True
+    config['compression']['params']['prune_last_conv'] = True
+    config['compression']['params']['all_weights'] = all_weights
+    input_shape = [1, 8, 8, 1]
+    model = model(input_shape)
+    model.compile()
+    _, compression_ctrl = create_compressed_model_and_algo_for_test(model, config)
+
+    assert compression_ctrl.full_flops == ref_full_flops
+    assert compression_ctrl.full_params_num == ref_full_params
+    assert compression_ctrl.current_flops == ref_current_flops
+    assert compression_ctrl.current_params_num == ref_current_params