[POT] Unify MatMul scales (#13717)

* Support FQ scales unification for MatMul

* Add test

* Update test_scales.py

tools/pot/openvino/tools/pot/algorithms/quantization/fake_quantize_configuration.py

@@ -12,6 +12,7 @@ from ...graph.utils import find_operation_matches, get_operation_list, is_data_t
 from ...graph.model_utils import get_nodes_by_type, get_node_by_name
 from ...graph.node_utils import get_input_shape, get_all_node_outputs,\
 get_node_input, get_node_inputs, get_node_data_type, check_const_input
+from ...graph.passes import traverse_graph
 
 from ...utils.logger import get_logger
 
@@ -309,7 +310,7 @@ def find_fqs_to_unify(model, config):
         return True
 
     def _is_concat_unify_condition(node):
-        def _is_followed_by_conv(input_node):
+        def _is_followed_by_concat_unify_outs_op(input_node):
             if _is_quantize_agnostic_op(input_node):
                 concat_stack.extend(get_all_node_outputs(input_node))
             elif input_node.type in [n['type'] for n in CONCAT_UNIFY_OUTPUTS]:
@@ -331,18 +332,23 @@ def find_fqs_to_unify(model, config):
                     return True
             return False
 
-        res = False
-        concat_inputs = get_node_inputs(node)
-        for concat_input in concat_inputs:
-            if concat_input.type not in [n['type'] for n in CONCAT_UNIFY_INPUTS]:
-                logger.debug('Concat %s without FQ or Concat as input will not unified',
-                             node.fullname)
-                return res
+        def _source_fn(op):
+            return [p for p in get_node_inputs(op) if p and p.type != 'Const']
+
+        def _criteria_fn(op):
+            return op.type in [n['type'] for n in CONCAT_UNIFY_INPUTS]
+
+        _, criteria = traverse_graph(node,
+                                    move_fn=_source_fn,
+                                    stop_criteria_fn=_criteria_fn,
+                                    criteria_fns=_criteria_fn)
+        if not criteria:
+            return False
 
         concat_stack = [node]
         while concat_stack:
             node_to_check = concat_stack.pop()
-            res = _is_followed_by_conv(node_to_check)
+            res = _is_followed_by_concat_unify_outs_op(node_to_check)
         if res:
             concat_stack = [node]
             while concat_stack:

tools/pot/openvino/tools/pot/graph/special_operations.py

@@ -42,7 +42,8 @@ OPERATIONS_WITH_WEIGHTS = [
 
 CONCAT_UNIFY_OUTPUTS = [
     {'type': 'ConvolutionBackpropData'},
-    {'type': 'Convolution'}
+    {'type': 'Convolution'},
+    {'type': 'MatMul'},
 ]
 
 CONCAT_UNIFY_INPUTS = [

tools/pot/tests/data/models/matmul_unify_scales/matmul_1_in_split.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:23079d4cb21aeef7d1db429126dfc3bde41b579ad21637c807959115e12330ed
+size 154

tools/pot/tests/data/models/matmul_unify_scales/matmul_1_in_split.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c55f22340157b9d2806dce28deb9b5fc4f20feaa326678e08d22199e3b762a88
+size 6056

tools/pot/tests/data/models/matmul_unify_scales/matmul_3_in_agnostic.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:203b4e8da5c653b197e2407e7497a7fee5cf1eebc7d7bfe8c88d5d32ba6d7e8d
+size 160

tools/pot/tests/data/models/matmul_unify_scales/matmul_3_in_agnostic.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5800791d0296272644bd090f274c53dcd67cec3ecf634e6f9e321d4ce27c713a
+size 1834

tools/pot/tests/data/models/matmul_unify_scales/matmul_3_in_cat.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f06423af514397433aa6fb69979b515b492464a7d59ad41e518a02517d77cedc
+size 150

tools/pot/tests/data/models/matmul_unify_scales/matmul_3_in_cat.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3ebe47925cce356f0b5739477c943124dee566996004e064f9baa957a705955a
+size 493

tools/pot/tests/data/models/matmul_unify_scales/matmul_3_in_slice.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:84214999c20b6ab060aeec75afe2f2b26ef29bbb8460c2df0578fd3d9b1a31a2
+size 154

tools/pot/tests/data/models/matmul_unify_scales/matmul_3_in_slice.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f9862637721fedde79f9ca30b555a7f693e7384a38bd318f9ee50640ab8ee64f
+size 1595

tools/pot/tests/data/reference_fake_quantize_conf/matmul_1_in_split.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:29c3f04406a2d0fcb069c73d32ae4b22a7158fc69a3d85f3e357d2b5d2217220
+size 398

tools/pot/tests/data/reference_fake_quantize_conf/matmul_3_in_agnostic.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f89b634b6d956b7fe04230d5260515fb00d816a62c6d4ebbde69e4ff1b479d72
+size 573

tools/pot/tests/data/reference_fake_quantize_conf/matmul_3_in_cat.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fe556eda19d735c0088b4453730019ff5c75a0bfb0e2b0ed19147d29c8e00498
+size 372

tools/pot/tests/data/reference_fake_quantize_conf/matmul_3_in_slice.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c2b3afcb32fdc37395bfd0fb55dc1684daf8c46d5af839e5ef745c861beb7da3
+size 571

tools/pot/tests/test_scales.py

@@ -9,8 +9,10 @@ import numpy as np
 import pytest
 from addict import Dict
 
+from openvino.tools.pot import DataLoader
 from openvino.tools.pot.algorithms.algorithm_selector import COMPRESSION_ALGORITHMS
 from openvino.tools.pot.engines.ac_engine import ACEngine
+from openvino.tools.pot.engines.simplified_engine import SimplifiedEngine
 from openvino.tools.pot.graph import load_model
 from openvino.tools.pot.graph.node_utils import get_node_inputs, get_node_input, get_node_value
 from openvino.tools.pot.graph import model_utils as mu
@@ -181,6 +183,58 @@ def load_refs(path_to_refs):
         return json.load(json_file)
 
 
+def make_list(x):
+    if isinstance(x, np.ndarray):
+        x = x.tolist()
+    if isinstance(x, list):
+        return x
+    return [x]
+
+
+def calculate_statistics(model, compression_config, engine):
+    algo = COMPRESSION_ALGORITHMS.get('MinMaxQuantization')(compression_config, engine)
+    stats_collector = StatisticsCollector(engine)
+    algo.register_statistics(model, stats_collector)
+    stats_collector.compute_statistics(model)
+    return algo.run(model)
+
+
+def aggregate_fq_statistics(fq_list):
+    model_values = {}
+    for fq in sorted(fq_list, key=lambda item: item.name):
+        min_levels, max_levels = tuple([get_node_value(node)
+                                        for node in get_node_inputs(fq)[1:3]])
+        fq_name = fq.fullname
+        if get_node_input(fq, 0).type == 'Const':
+            min_levels = min_levels.reshape(min_levels.shape[0])
+            max_levels = max_levels.reshape(max_levels.shape[0])
+        else:
+            if not min_levels.shape and not max_levels.shape:
+                pass
+            else:
+                min_levels = min_levels.reshape(min_levels.shape[1])
+                max_levels = max_levels.reshape(max_levels.shape[1])
+
+        print(fq_name, 'max_levels', max_levels)
+        min_levels = make_list(min_levels)
+        max_levels = make_list(max_levels)
+        model_values[fq_name] = {'max': max_levels, 'min': min_levels}
+    return model_values
+
+
+def compare_min_max_values(model_values, refs):
+    for ref_name in refs:
+        refs_min_levels = make_list(refs[ref_name]['min'])
+        refs_max_levels = make_list(refs[ref_name]['max'])
+        min_levels = model_values[ref_name]['min']
+        max_levels = model_values[ref_name]['max']
+
+        for min_level, max_level, ref_min, ref_max in zip(
+                min_levels, max_levels, refs_min_levels, refs_max_levels):
+            assert abs(min_level - ref_min) < EPS
+            assert abs(max_level - ref_max) < EPS
+
+
 CONFIGURATIONS = [
     ('mobilenet-v2-pytorch', 'pytorch', 'symmetric'),
     ('mobilenet-v2-pytorch', 'pytorch', 'asymmetric'),
@@ -233,69 +287,92 @@ def test_fake_quantize_configurations(tmp_path, models, model_name, model_framew
         }
     })
 
-    def _make_list(x):
-        if isinstance(x, np.ndarray):
-            x = x.tolist()
-        if isinstance(x, list):
-            return x
-        return [x]
-
     engine = ACEngine(config)
     compression_config.subset_indices = [0]
-    algo = COMPRESSION_ALGORITHMS.get('MinMaxQuantization')(compression_config, engine)
     model = models.get(model_name, model_framework, tmp_path)
     model = load_model(model.model_params)
 
-    stats_collector = StatisticsCollector(engine)
-    algo.register_statistics(model, stats_collector)
-    stats_collector.compute_statistics(model)
-
-    model = algo.run(model)
+    model_with_stats = calculate_statistics(model, compression_config, engine)
+    fq_list = mu.get_nodes_by_type(model_with_stats, ['FakeQuantize'])
+    model_values = aggregate_fq_statistics(fq_list)
 
     refs_path = os.path.join(REFERENCES_DIR, '{}_{}.json'.format(model_name, algo_mode))
     local_path = os.path.join(tmp_path, '{}.json'.format(model_name))
-
     ref_exists = os.path.isfile(refs_path)
-
     refs = load_refs(refs_path) if ref_exists else {}
-    ref_file = None if ref_exists else open(refs_path, 'w')
     local_file = open(local_path, 'w')
-    model_values = {}
-
-    fq_list = mu.get_nodes_by_type(model, ['FakeQuantize'])
-    for fq in sorted(fq_list, key=lambda item: item.name):
-        min_levels, max_levels = tuple([get_node_value(node)
-                                        for node in get_node_inputs(fq)[1:3]])
-        fq_name = fq.fullname
-        if get_node_input(fq, 0).type == 'Const':
-            min_levels = min_levels.reshape(min_levels.shape[0])
-            max_levels = max_levels.reshape(max_levels.shape[0])
-        else:
-            if not min_levels.shape and not max_levels.shape:
-                pass
-            else:
-                min_levels = min_levels.reshape(min_levels.shape[1])
-                max_levels = max_levels.reshape(max_levels.shape[1])
-
-        min_levels = _make_list(min_levels)
-        max_levels = _make_list(max_levels)
-        model_values[fq_name] = {'max': max_levels, 'min': min_levels}
 
     if not ref_exists:
+        ref_file = open(refs_path, 'w')
         json.dump(model_values, ref_file)
         return
     json.dump(model_values, local_file)
 
-    for ref_name in refs:
-        refs_min_levels = _make_list(refs[ref_name]['min'])
-        refs_max_levels = _make_list(refs[ref_name]['max'])
-        min_levels = model_values[ref_name]['min']
-        max_levels = model_values[ref_name]['max']
+    compare_min_max_values(model_values, refs)
 
-        for min_level, max_level, ref_min, ref_max in zip(
-                min_levels, max_levels, refs_min_levels, refs_max_levels):
-            assert abs(min_level - ref_min) < EPS
-            assert abs(max_level - ref_max) < EPS
+
+MATMUL_MODELS = [
+    ('matmul_3_in_slice', 'pytorch',
+     {'onnx::Concat_0': [5, 10], 'onnx::Concat_1': [5, 2], 'onnx::Concat_2': [5, 10]}),
+    ('matmul_3_in_cat', 'pytorch',
+     {'onnx::Concat_0': [5, 10], 'onnx::Concat_1': [5, 2], 'onnx::Concat_2': [5, 10]}),
+    ('matmul_3_in_agnostic', 'pytorch',
+     {'onnx::Concat_0': [5, 10], 'onnx::Concat_1': [5, 2], 'onnx::Transpose_2': [5, 10]}),
+    ('matmul_1_in_split', 'pytorch',
+     {'onnx::Squeeze_0': [1, 3, 224, 224]}),
+]
+
+
+class RandomDataLoader(DataLoader):
+    """ Generate random float32 data """
+    def __init__(self, shapes, seed=0):
+        self.rng = np.random.default_rng(seed)
+        data = {}
+        for i, (name, shape) in enumerate(shapes.items()):
+            data[name] = self.rng.uniform(-1 + 5 * i, 1 + 5 * i, shape)
+        self._data = [data]
+
+    def __len__(self):
+        return 1
+
+    def __getitem__(self, index):
+        return self._data[index]
+
+
+@pytest.mark.parametrize(
+    'model_name, model_framework, input_shapes', MATMUL_MODELS,
+    ids=['{}_{}'.format(m[0], m[1]) for m in MATMUL_MODELS]
+)
+def test_matmul_scale_unification(tmp_path, models, model_name, model_framework, input_shapes):
+    model = models.get(model_name, model_framework, tmp_path)
+    model = load_model(model.model_params)
+
+    compression_config = Dict({
+        'name': 'DefaultQuantization',
+        'stat_subset_size': 1,
+        'target_device': 'CPU',
+    })
+
+    data_loader = RandomDataLoader(input_shapes)
+    engine = SimplifiedEngine({"device": "CPU"}, data_loader=data_loader)
+    model_with_stats = calculate_statistics(model, compression_config, engine)
+
+    fq_list = mu.get_nodes_by_type(model_with_stats, ['FakeQuantize'])
+    model_values = aggregate_fq_statistics(fq_list)
+
+    refs_path = os.path.join(REFERENCES_DIR, f'{model_name}.json')
+    local_path = os.path.join(tmp_path, f'{model_name}.json')
+    ref_exists = os.path.isfile(refs_path)
+    refs = load_refs(refs_path) if ref_exists else {}
+    local_file = open(local_path, 'w')
+
+    if not ref_exists:
+        ref_file = open(refs_path, 'w')
+        json.dump(model_values, ref_file)
+        return
+    json.dump(model_values, local_file)
+
+    compare_min_max_values(model_values, refs)
 
 
 def _get_pytorch_accuracy_checker_config(path_to_dataset):