[POT] Update model compression pass and FQ values precision (#8855)

* Update model compression pass and FQ values precision

* Pylint fixes

* Fix for case with 2 or more destinations

* Update pass.py from suggestion

Co-authored-by: Anton Chetverikov <Anton.Chetverikov@intel.com>

* Update with the const condition

* Update ref

* Revert ref and update data_type for FQ consts

Co-authored-by: Anton Chetverikov <Anton.Chetverikov@intel.com>

tools/pot/openvino/tools/pot/algorithms/quantization/fast_bias_correction/algorithm.py

@@ -49,6 +49,7 @@ class FastBiasCorrection(Algorithm):
          :param model: model to apply algo
          :return with corrected biases for layers with bias
          """
+        mu.nx_type_infer(model)
         activations_statistics = self._stats_collector.get_statistics_for_algorithm(self.name)
         nodes_with_bias = mu.get_nodes_by_type(model, [op['type'] for op in OPERATIONS_WITH_BIAS])
         self.find_channel_axis(model)

tools/pot/openvino/tools/pot/graph/builder.py

@@ -2,11 +2,10 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from copy import deepcopy
-import numpy as np
 
 from mo.graph.graph import Graph
 
-from openvino.tools.pot.graph.node_utils import get_node_input, get_node_inputs
+from openvino.tools.pot.graph.node_utils import get_node_data_type, get_node_input, get_node_inputs
 from .editor import create_node, connect_nodes_by_name, get_node_by_name
 
 
@@ -105,8 +104,10 @@ def build_graph_for_node(model, input_name, input_shape, node, remove_bias=False
      :param remove_fake_quantize: remove fake quantize nodes in the generated graph
      :return: generated graph.
     """
+    input_data_type = get_node_data_type(node, 0)
     nodes, edges = [], []
-    nodes.append((input_name, 'Parameter', {'name': input_name, 'shape': input_shape, 'type': 'Parameter'}))
+    nodes.append((input_name, 'Parameter', {'name': input_name, 'shape': input_shape,
+                                            'type': 'Parameter', 'data_type': input_data_type}))
 
     node_attrs = deepcopy(node.attrs())
     if node.has_valid('output') and node.has_valid('get_output_feature_dim'):
@@ -146,6 +147,9 @@ def build_graph_for_node(model, input_name, input_shape, node, remove_bias=False
     weights_node = get_node_input(src_node, 1)
     weights_node = get_node_input(weights_node, 0) \
         if weights_node.type == 'FakeQuantize' else weights_node
-    if weights_node.out_port(0).get_data_type() == np.float16:
+    weights_out_dtype = weights_node.out_port(0).get_data_type()
+    src_out_dtype = src_node.out_port(0).get_data_type()
+    if weights_out_dtype != src_out_dtype:
         weights_node.out_node(0)['Insert_Convert_operation_after'] = True
+
     return graph

tools/pot/openvino/tools/pot/graph/model_utils.py

@@ -3,6 +3,7 @@
 
 import networkx as nx
 from mo.graph.graph import Node
+from mo.middle.passes.infer import type_infer
 
 from . import editor as ge, builder as gb
 from .nx_model import NXModel
@@ -127,3 +128,9 @@ def models_union(first_model, second_model):
         union_dict['model'].graph.update(model_1.graph)
         union_dict['model'].graph.update(model_2.graph)
     return union
+
+def nx_type_infer(model):
+    """ Apply type_infer for each model in NXModel wrapper
+    """
+    for model_dict in model.models:
+        type_infer(model_dict['model'])

tools/pot/openvino/tools/pot/graph/node_utils.py

@@ -267,10 +267,10 @@ def get_lstm_ends(read_value, assigns, ignore_nodes):
     return lstm_outputs
 
 
-def get_node_data_type(node):
-    if node.type != 'Const' and node.in_port(0).get_source() is not None \
-            and node.in_port(0).get_source().is_data_type_defined():
-        return node.in_port(0).get_source().get_data_type()
+def get_node_data_type(node, port_id=0):
+    if node.type != 'Const' and node.in_port(port_id).get_source() is not None \
+            and node.in_port(port_id).get_source().is_data_type_defined():
+        return node.in_port(port_id).get_source().get_data_type()
     return None
 
 

tools/pot/openvino/tools/pot/graph/passes.py

@@ -22,7 +22,6 @@ from mo.graph.port import Port
 from mo.middle.pattern_match import apply_pattern
 from mo.ops.const import Const
 from mo.middle.passes.convert_data_type import convert_blob
-from mo.middle.passes.infer import type_infer
 
 from . import editor as ge
 from . import node_utils as nu
@@ -221,7 +220,6 @@ class FakeQuantizePropagation(BackReplacementPattern):
     }
 
     def delete_fq_non_quantizable_node_precision(self, graph):
-        type_infer(graph)
         fq_removal = RemoveFakeQuantize()
         fq_removal.quantize_agnostic_operations = self.quantize_agnostic_operations
         fq_removal.quantize_operations = self.quantize_operations
@@ -713,14 +711,14 @@ def create_bias_node(graph: Graph, src_node):
     add_bias.out_node(0)['Insert_Convert_operation_after'] = True
 
 
-def create_fake_quantize_node(graph: Graph, name):
+def create_fake_quantize_node(graph: Graph, name, data_type=np.float32):
     fq = FakeQuantize(graph, {'name': name, 'levels': 0,
                               'stop_value_propagation': True}).create_node()
 
-    input_low = Const(graph, {'value': np.array(0.0).astype(np.float32)}).create_node()
-    input_height = Const(graph, {'value': np.array(0.0).astype(np.float32)}).create_node()
-    output_low = Const(graph, {'value': np.array(0.0).astype(np.float32)}).create_node()
-    output_height = Const(graph, {'value': np.array(0.0).astype(np.float32)}).create_node()
+    input_low = Const(graph, {'value': np.array(0.0, dtype=data_type)}).create_node()
+    input_height = Const(graph, {'value': np.array(0.0, dtype=data_type)}).create_node()
+    output_low = Const(graph, {'value': np.array(0.0, dtype=data_type)}).create_node()
+    output_height = Const(graph, {'value': np.array(0.0, dtype=data_type)}).create_node()
 
     input_low.out_port(0).connect(fq.in_port(1))
     input_height.out_port(0).connect(fq.in_port(2))
@@ -764,9 +762,11 @@ def insert_fake_quantize(graph, node, ports=None, names=None):
         if port_name is not None and idx in port_name:
             name = port_name[idx]
 
+        port_data_type = nu.get_node_data_type(node, idx)
+        port_data_type = port_data_type if port_data_type else np.float32
         # Create FakeQuantize operations
         fq_input = create_fake_quantize_node(
-            graph, '{node_name}/{name}_{idx}'.format(node_name=node.name, name=name, idx=idx))
+            graph, '{node_name}/{name}_{idx}'.format(node_name=node.name, name=name, idx=idx), port_data_type)
 
         # Insert FakeQuantize after input
         if node.type == 'Result':
@@ -841,6 +841,7 @@ def traverse_graph(node, move_fn, stop_criteria_fn=None, criteria_fns=None):
 
 def compress_weights(model: Graph):
     """Apply transformations to save model weights to INT8."""
+    add_removed_converts(model)
     CompressQuantizeWeights().find_and_replace_pattern(model)
     model.clean_up()
     ForceStrictPrecision().find_and_replace_pattern(model)
@@ -906,18 +907,25 @@ def add_removed_converts(graph: Graph):
         data_node = Node(graph, data_node_name)
         # Get access to Const node connected to data node
         const_op = data_node.in_node(0)
-        assert const_op.data_type == np.float32, "Error when try to insert Convert operation after Const: {}".\
-            format(const_op.soft_get('name'))
+
+        if const_op.type != 'Const':
+            logger.debug('Error when try to insert Convert operation after {} with {} type'.\
+                format(const_op.soft_get('name'), const_op.soft_get('type')))
+            continue
+
+        if const_op.data_type != np.float32:
+            logger.debug('Error when try to insert Convert operation after Const: {}'.\
+                format(const_op.soft_get('name')))
+            continue
 
         convert_op = Cast(graph, {'dst_type': np.float32,
                                   'name': const_op.name + '/restored_convert',
                                   'stop_value_propagation': True}).create_node()
 
         # Insert Convert operation after Const operation
-        consumer_port = const_op.out_port(0).get_connection().get_destination()
-        const_op.out_port(0).get_connection().set_destination(convert_op.in_port(0))
-        convert_op.out_port(0).connect(consumer_port)
+        const_op.out_port(0).get_connection().insert_node(convert_op)
+        convert_op.out_node().value = None
 
-        # Convert Const value to FP32 to make types in graph consistent
+        # Convert Const value to FP16 to make types in graph consistent
         const_op.value, _, _ = convert_blob(const_op.value, np.float16)
         const_op.infer(const_op)

tools/pot/openvino/tools/pot/graph/transformer.py

@@ -1,6 +1,8 @@
 # Copyright (C) 2020-2021 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
+from mo.middle.passes.infer import type_infer
+
 from .editor import add_fullname_for_nodes
 from .special_operations import QUANTIZE_AGNOSTIC_OPERATIONS
 from .passes import InsertFakeQuantize, FakeQuantizePropagation, FakeQuantizeOptimization, RemoveFakeQuantize, \
@@ -46,6 +48,8 @@ class GraphTransformer:
         self.nodes_marker.mark_ignored_blocks(graph, self.target_device)
         graph.clean_up()
 
+        type_infer(graph)
+
         self.fq_insertion.find_and_replace_pattern(graph)
         graph.clean_up()