Reverse infer (#8767)

model-optimizer/extensions/front/tf/placeholder_ext.py

@@ -18,5 +18,7 @@ class PlaceholderFrontExtractor(FrontExtractorOp):
             'shape': tf_tensor_shape(node.pb.attr["shape"].shape),
             'permute_attrs': PermuteAttrs().update_attrs(attrs=[('shape', 'output:0')])
         }
+        if node.pb.attr["shape"].shape.unknown_rank:
+            attrs['shape'] = None
         Parameter.update_node_stat(node, attrs)
         return cls.enabled

model-optimizer/extensions/ops/parameter.py

@@ -4,7 +4,7 @@
 import numpy as np
 
 from mo.front.common.partial_infer.utils import unmask_shape
-from mo.graph.graph import Graph
+from mo.graph.graph import Graph, Node
 from mo.middle.passes.convert_data_type import np_data_type_to_destination_type
 from mo.ops.op import Op, PermuteAttrs
 
@@ -19,6 +19,7 @@ class Parameter(Op):
             'version': 'opset1',
 
             'infer': self.infer,
+            'reverse_infer': self.reverse_infer,
             'is_input': True,
             'data_type': None,
 
@@ -49,3 +50,12 @@ class Parameter(Op):
         node.out_port(0).data.set_shape(node.shape)
 
         PermuteAttrs.create_permute_attrs(node, attrs=[('shape', 'output:0')])
+
+    @staticmethod
+    def reverse_infer(node: Node):
+        # update node 'shape' attribute (if it is not defined) from the output port shape which was calculated
+        # during the reverse_infer phase
+        shape = node.soft_get('shape', None)
+        if shape is None and node.out_port(0).data.get_shape() is not None:
+            node['shape'] = node.out_port(0).data.get_shape()
+

model-optimizer/mo/front/common/partial_infer/utils.py

@@ -24,6 +24,16 @@ def shape_array(value, dtype=np.int64):
     return np.ma.masked_equal(value, dynamic_dimension_value).astype(dtype=dtype)
 
 
+def undefined_shape_of_rank(rank: int):
+    """
+    Create a shape of specified rank with all dynamic dimensions.
+
+    :param rank: requested rank of the output shape
+    :return: shape array
+    """
+    return shape_array([dynamic_dimension_value] * rank)
+
+
 def compatible_dims(dim1, dim2):
     """
     Compare if dim1 is equal to dim2 or any of them is dynamic

model-optimizer/mo/middle/passes/infer.py

@@ -2,12 +2,12 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import logging as log
+from typing import List
 
 import networkx as nx
 
 from mo.front.common.partial_infer.utils import dynamic_dimension
-from mo.graph.graph import Node, Graph
-from mo.graph.graph import dict_includes
+from mo.graph.graph import Node, Graph, dict_includes
 from mo.utils.error import Error
 from mo.utils.utils import refer_to_faq_msg, shrink_str_value
 
@@ -93,18 +93,59 @@ def partial_infer(graph: Graph, start_node: str = None):
         nx.set_node_attributes(G=graph.subgraph(nodes[start_index:]), name='is_partial_inferred', values=False)
     else:
         nx.set_node_attributes(G=graph, name='is_partial_inferred', values=False)
-    debug_logger = log.getLogger().isEnabledFor(log.DEBUG)
 
     nx.set_node_attributes(G=graph, name='executable',
                            values={n: True for n in graph.get_nodes_with_attributes(kind='data')})
 
+    # first we infer constant sub-graphs so the reverse infer could use constant values sub-graphs. For example,
+    # convolution weights may be reshuffled by some operation in the graph and are not directly consumed by the conv
+    # node
+    infer_nodes(graph, nodes, True)
+
+    # we may need to deduce shape for Parameter node(s) if it is not defined
+    need_reverse_infer = False
+    for parameter in graph.get_op_nodes(op='Parameter'):
+        if parameter.soft_get('shape', None) is None:
+            need_reverse_infer = True
+
+    if need_reverse_infer:
+        reverse_infer(graph, nodes)
+
+    infer_nodes(graph, nodes, False)
+
+    not_fully_inferred = graph.get_nodes_with_attributes(is_not_fully_inferred=True)
+    for n in not_fully_inferred:
+        node = Node(graph, n)
+        if node.has('infer') and not node.infer is None:
+            node.infer(node)
+
+    return graph
+
+
+def infer_nodes(graph: Graph, nodes: List[Node], constant_subgraph_only: bool = False):
+    """
+    Run "infer" function of the specified nodes.
+
+    :param graph: graph with nodes
+    :param nodes: list of node ids in the topological order
+    :param constant_subgraph_only: flag which specifies whether only inference of constant sub-graphs should be done
+    """
+    debug_logger = log.getLogger().isEnabledFor(log.DEBUG)
     for n in nodes:
         # Data Flow Infer
+        node = Node(graph, n)
+        node_name = node.soft_get('name', node.id)
         try:
-            node = Node(graph, n)
-            node_name = node.soft_get('name')
             if node.has('is_partial_inferred') and not node.is_partial_inferred:
                 if node.has('infer') and not node.infer is None:
+                    # we consider that operation will produce value if all inputs are constants or it is
+                    # 'ShapeOf' operation
+                    if constant_subgraph_only:
+                        in_values = [port.data.get_value() for port in node.in_ports().values()]
+                        if node.soft_get('op') == 'Parameter' or any(value is None for value in in_values) or \
+                                (node.soft_get('op') == 'ShapeOf' and node.in_port(0).data.get_shape() is None):
+                            continue
+
                     if debug_logger:
                         log.debug('-' * 20)
                         log.debug('Partial infer for {}'.format(node.soft_get('name')))
@@ -125,18 +166,19 @@ def partial_infer(graph: Graph, start_node: str = None):
                         log.debug('Outputs:')
                         log_debug_dict(node.out_nodes(), 'output')
 
-                    not_all_output_shapes = False
-
-                    for out_port, out_node in out_nodes.items():
+                    if not constant_subgraph_only:
                         not_all_output_shapes = False
-                        if not out_node.has_valid('shape'):
-                            log.error('Shape is not defined for output {} of "{}".'.format(out_port, node_name))
-                            not_all_output_shapes = True
 
-                    if not_all_output_shapes:
-                        raise Error('Not all output shapes were inferred or fully defined for node "{}". ' +
-                                    refer_to_faq_msg(40),
-                                    node_name)
+                        for out_port, out_node in out_nodes.items():
+                            not_all_output_shapes = False
+                            if not out_node.has_valid('shape'):
+                                log.error('Shape is not defined for output {} of "{}".'.format(out_port, node_name))
+                                not_all_output_shapes = True
+
+                        if not_all_output_shapes:
+                            raise Error('Not all output shapes were inferred or fully defined for node "{}". ' +
+                                        refer_to_faq_msg(40),
+                                        node_name)
                 elif node.kind != 'data':
                     raise Error(
                         'There is no registered "infer" function for node "{}" with op = "{}". ' +
@@ -146,7 +188,6 @@ def partial_infer(graph: Graph, start_node: str = None):
                         node.soft_get('op')
                     )
                 node.is_partial_inferred = True
-
         except Exception as err:
             log.error('Cannot infer shapes or values for node "{}".'.format(node.soft_get('name')))
             log.error(str(err))
@@ -169,14 +210,6 @@ def partial_infer(graph: Graph, start_node: str = None):
                         refer_to_faq_msg(38)) from err
         control_flow_infer(graph, n)
 
-    not_fully_inferred = graph.get_nodes_with_attributes(is_not_fully_inferred=True)
-    for n in not_fully_inferred:
-        node = Node(graph, n)
-        if node.has('infer') and not node.infer is None:
-            node.infer(node)
-
-    return graph
-
 
 def override_batch(graph: Graph, batch: int):
     """
@@ -286,3 +319,13 @@ def copy_type_infer(node):
                 out_port.set_data_type(connected_in_ports[0].get_data_type())
         else:
             raise Error('No input ports of node {} to determine data type'.format(node.soft_get('name')))
+
+
+def reverse_infer(graph: Graph, nodes: list):
+    nodes = reversed(nodes)
+    for n in nodes:
+        node = Node(graph, n)
+        if node.has_valid('reverse_infer'):
+            log.debug("Executed reverse infer for node '{}'".format(node.soft_get('name', node.id)))
+            node.reverse_infer(node)
+ 

model-optimizer/mo/ops/convolution.py

@@ -6,7 +6,7 @@ import logging as log
 import numpy as np
 
 from mo.front.common.partial_infer.utils import int64_array, mark_input_bins, assign_dims_to_weights, \
-    tf_window_op_pad_infer, dynamic_dimension_value, shape_array
+    tf_window_op_pad_infer, dynamic_dimension_value, shape_array, is_fully_defined, undefined_shape_of_rank
 from mo.front.onnx.extractors.utils import get_backend_pad
 from mo.graph.graph import Node, Graph
 from mo.graph.perm_inputs import PermuteInputs
@@ -23,6 +23,7 @@ class Convolution(Op):
             'op': self.op,
             'version': 'opset1',
             'infer': self.infer,
+            'reverse_infer': self.reverse_infer,
             'multiplication_transparent': True,
             'multiplication_transparent_ports': [(0, 0), (1, 0)],
             'in_ports_count': 3,
@@ -112,18 +113,29 @@ class Convolution(Op):
                 log.error('Cannot reshape kernel due to not all required attrs was set to {} node'.format(node.id))
                 return
             # layout for Convolution weights is OIHW
-            kernel_shape = int64_array([node.output, input_shape[node.channel_dims].item() / node.group,
+            kernel_shape = shape_array([node.output, input_shape[node.channel_dims].item() / node.group,
                                        *[node.kernel_spatial[i] for i in range(len(node.kernel_spatial))]])
             if node.type == 'Deconvolution':  # layout for Deconvolution weights is IOHW
                 kernel_shape[[0, 1]] = kernel_shape[[1, 0]]
 
-            if np.prod(kernel_shape) != np.prod(node.in_node(weights_index).value.shape):
+            if is_fully_defined(kernel_shape) and np.prod(kernel_shape) != np.prod(node.in_node(weights_index).value.shape):
                 log.error("Size of weights {} does not match kernel shape: {}\n"
                           "".format(np.prod(node.in_node(weights_index).value.shape), kernel_shape) +
                           "    Possible reason is wrong channel number in input shape\n")
                 raise Error("Cannot reshape weights to kernel shape")
 
-            node.in_node(weights_index).shape = np.array(kernel_shape)
+            if not is_fully_defined(kernel_shape):
+                num_undefined = np.count_nonzero(kernel_shape.mask is True)  # pylint: disable=no-member
+                if num_undefined > 1:
+                    raise Error('Too many undefined dimensions of the kernel shape for node {}. Use --input_shape '
+                                'command line parameter to specify model input shapes'.format(node.soft_get('name',
+                                                                                                            node.id)))
+                kernel_size = np.prod(node.in_node(weights_index).value.shape)
+                # calculate undefined dimension using fully defined shape of the weights input and known kernel_shape
+                # dimensions
+                kernel_shape[np.where(kernel_shape == np.ma.masked)[0][0]] = kernel_size // np.prod(kernel_shape)
+
+            node.in_node(weights_index).shape = shape_array(kernel_shape)
             node.in_node(weights_index).value = np.reshape(node.in_node(weights_index).value, kernel_shape)
             node.reshape_kernel = False
 
@@ -262,3 +274,16 @@ class Convolution(Op):
         PermuteAttrs.set_permutation(node.in_node(weights_index), node, node.soft_get('get_weights_permute', None))
         PermuteInputs().set_input_permutation(
             node.in_node(weights_index), node, 'input:{}'.format(weights_index), 'transpose')
+
+    @staticmethod
+    def reverse_infer(node: Node):
+        input_shape = node.in_port(0).data.get_shape()
+        if input_shape is None:
+            shape = None
+            # TODO FIXME this is ugly solution based on various attributes which may not be set in some cases
+            for attr in ['dilation', 'stride', 'pad']:
+                if node.has_valid(attr):
+                    shape = undefined_shape_of_rank(len(node.soft_get(attr)))
+                    break
+            if shape is not None:
+                node.in_port(0).data.set_shape(shape)

model-optimizer/mo/ops/pad.py

@@ -3,8 +3,8 @@
 
 import numpy as np
 
-from mo.front.common.partial_infer.utils import is_fully_defined, shape_array
-from mo.graph.graph import Graph
+from mo.front.common.partial_infer.utils import is_fully_defined, shape_array, undefined_shape_of_rank
+from mo.graph.graph import Graph, Node
 from mo.graph.perm_inputs import PermuteInputs
 from mo.ops.op import Op
 
@@ -29,6 +29,7 @@ class Pad(Op):
 
             'version': 'opset1',
             'infer': self.infer,
+            'reverse_infer': self.reverse_infer,
 
             'mode': 'constant',
 
@@ -85,6 +86,13 @@ class Pad(Op):
         PermuteInputs().set_input_permutation(node.in_node(1), node, 'input:0', 'shape')
         PermuteInputs().set_input_permutation(node.in_node(2), node, 'input:0', 'shape')
 
+    @staticmethod
+    def reverse_infer(node: Node):
+        input_shape = node.in_port(0).data.get_shape()
+        if input_shape is None and node.is_in_port_connected(2) and node.in_port(2).data.get_shape() is not None:
+            shape = undefined_shape_of_rank(node.in_port(2).data.get_shape()[0])
+            node.in_port(0).data.set_shape(shape)
+
 
 class AttributedPad(Op):
     """ Pad operation that explicitly extends an input tensor at borders.
@@ -131,6 +139,7 @@ class TFPad(Op):
             'mode': 'constant',
         }, attrs)
 
+
 class ONNXPad(Op):
     """ Pad operation that explicitly extends an input tensor at borders.
 

model-optimizer/mo/ops/pooling.py

@@ -4,7 +4,7 @@
 import numpy as np
 
 from mo.front.common.partial_infer.utils import tf_window_op_pad_infer, int64_array, shape_array, \
-    dynamic_dimension_value, dynamic_dimension
+    dynamic_dimension_value, dynamic_dimension, undefined_shape_of_rank
 from mo.front.onnx.extractors.utils import get_backend_pad
 from mo.graph.graph import Node, Graph
 from mo.middle.passes.convert_data_type import np_data_type_to_destination_type
@@ -13,6 +13,12 @@ from mo.utils.error import Error
 from mo.front.extractor import bool_to_str
 
 
+poolings_map = {
+    'max': {'version': 'opset8', 'out_ports_count': 2},
+    'avg': {'version': 'opset1', 'out_ports_count': 1}
+}
+
+
 class PoolingV2(Op):
     """
     TensorFlow MaxPoolV2 and AvgPoolV2 operations expect windows_size and strides values from inputs not from
@@ -28,29 +34,35 @@ class PoolingV2(Op):
             'op': self.op,
             'version': None,
             'infer': self.infer,
+            'reverse_infer': self.reverse_infer,
             'in_ports_count': 3,
             'out_ports_count': 1,
         }, attrs)
 
     @staticmethod
     def infer(node: Node):
-        assert (len(node.in_nodes()) == 3), 'MaxPoolV2 node {} from must have only 3 inputs: input, window size, and strides ' \
-                                            'but instead got {} inputs'.format(node.soft_get('name', node.id), len(node.in_nodes()))
+        assert (len(node.in_nodes()) == 3), 'MaxPoolV2 node {} from must have only 3 inputs: input, window size, and ' \
+                                            'strides but instead got {} inputs'.format(node.soft_get('name', node.id),
+                                                                                       len(node.in_nodes()))
         node['window'] = node.in_port(1).data.get_value()
         node['stride'] = node.in_port(2).data.get_value()
 
         if node['window'] is None:
-            raise Error('The non-constant window size for MaxPoolV2 node {} is not supported'.format(node.soft_get('name', node.id)))
+            raise Error('The non-constant window size for MaxPoolV2 node {} is not supported'
+                        ''.format(node.soft_get('name', node.id)))
         if node['stride'] is None:
-            raise Error('The non-constant strides for MaxPoolV2 node {} is not supported'.format(node.soft_get('name', node.id)))
+            raise Error('The non-constant strides for MaxPoolV2 node {} is not supported'
+                        ''.format(node.soft_get('name', node.id)))
 
         Pooling.pool_infer(node)
 
-
-poolings_map = {
-    'max': {'version': 'opset8', 'out_ports_count': 2},
-    'avg': {'version': 'opset1', 'out_ports_count': 1}
-}
+    @staticmethod
+    def reverse_infer(node: Node):
+        input_shape = node.in_port(0).data.get_shape()
+        window_shape = node.in_port(1).data.get_shape()
+        # use the value of the 'window' input to determine input tensor rank
+        if input_shape is None and window_shape is not None:
+            node.in_port(0).data.set_shape(undefined_shape_of_rank(window_shape[0]))
 
 
 class Pooling(Op):
@@ -62,8 +74,10 @@ class Pooling(Op):
             'op': self.op,
             'version': poolings_map[attrs.get('pool_method')]['version'],
             'infer': self.infer,
+            'reverse_infer': self.reverse_infer,
             'in_ports_count': 1,
-            'out_ports_count': 1 if attrs.get('version') == 'opset1' else poolings_map[attrs.get('pool_method')]['out_ports_count']
+            'out_ports_count': 1 if attrs.get('version') == 'opset1' else
+            poolings_map[attrs.get('pool_method')]['out_ports_count']
         }, attrs)
 
     def backend_attrs(self):
@@ -184,3 +198,10 @@ class Pooling(Op):
                                                        ('window', 'input:0'),
                                                        ('spatial_dims', 'input:0'),
                                                        ('dilation', 'input:0')])
+
+    @staticmethod
+    def reverse_infer(node: Node):
+        input_shape = node.in_port(0).data.get_shape()
+        window = node.soft_get('window', None)
+        if input_shape is None and window is not None:
+            node.in_port(0).data.set_shape(undefined_shape_of_rank(len(window)))

model-optimizer/mo/ops/squeeze.py

@@ -4,7 +4,8 @@
 import numpy as np
 
 from mo.front.caffe.extractors.utils import get_canonical_axis_index
-from mo.front.common.partial_infer.utils import int64_array, dynamic_dimension, shape_delete, is_fully_defined
+from mo.front.common.partial_infer.utils import int64_array, dynamic_dimension, shape_delete, is_fully_defined, \
+    undefined_shape_of_rank
 from mo.graph.graph import Node
 from mo.graph.perm_inputs import PermuteInputs
 from mo.ops.op import Op
@@ -26,6 +27,7 @@ class Squeeze(Op):
             'in_ports_count': 2,
             'out_ports_count': 1,
             'infer': self.infer,
+            'reverse_infer': self.reverse_infer,
         }, attrs)
 
     @staticmethod
@@ -69,3 +71,13 @@ class Squeeze(Op):
 
         # the squeeze_dim attribute will be converted to the second input in the end of the Middle phase
         PermuteInputs().set_input_permutation(node.in_node(1), node, 'input:0', 'axis')
+
+    @staticmethod
+    def reverse_infer(node: Node):
+        input_shape = node.in_port(0).data.get_shape()
+        output_shape = node.out_port(0).data.get_shape()
+        squeeze_dims = node.in_port(1).data.get_value()
+        if input_shape is None and output_shape is not None and squeeze_dims is not None:
+            num_squeeze_dims = 1 if int64_array(squeeze_dims).ndim == 0 else len(squeeze_dims)
+            shape = undefined_shape_of_rank(len(output_shape) + num_squeeze_dims)
+            node.in_port(0).data.set_shape(shape)

model-optimizer/mo/ops/unsqueeze.py

@@ -1,9 +1,8 @@
 # Copyright (C) 2018-2021 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-import numpy as np
-
-from mo.front.common.partial_infer.utils import int64_array, shape_array, is_fully_defined, shape_insert
+from mo.front.common.partial_infer.utils import int64_array, is_fully_defined, shape_insert, undefined_shape_of_rank
+from mo.graph.graph import Node
 from mo.graph.perm_inputs import PermuteInputs
 from mo.ops.op import Op
 from mo.utils.error import Error
@@ -26,7 +25,8 @@ class Unsqueeze(Op):
             'reinterp_shape': True,
             'in_ports_count': 2,
             'out_ports_count': 1,
-            'infer': self.infer
+            'infer': self.infer,
+            'reverse_infer': self.reverse_infer,
         }, attrs)
 
     @staticmethod
@@ -61,3 +61,13 @@ class Unsqueeze(Op):
             node.out_port(0).data.set_shape(output_shape)
 
         PermuteInputs().set_input_permutation(node.in_node(1), node, 'input:0', 'axis')
+
+    @staticmethod
+    def reverse_infer(node: Node):
+        input_shape = node.in_port(0).data.get_shape()
+        output_shape = node.out_port(0).data.get_shape()
+        unsqueeze_dims = node.in_port(1).data.get_value()
+        if input_shape is None and output_shape is not None and unsqueeze_dims is not None:
+            num_unsqueeze_dims = 1 if int64_array(unsqueeze_dims).ndim == 0 else len(unsqueeze_dims)
+            shape = undefined_shape_of_rank(len(output_shape) - num_unsqueeze_dims)
+            node.in_port(0).data.set_shape(shape)

model-optimizer/mo/utils/telemetry_utils.py

@@ -54,7 +54,7 @@ def send_shapes_info(framework: str, graph: Graph):
         shape_str = ""
         is_partially_defined = "0"
         for shape in shapes:
-            shape_str += np.array2string(int64_array(unmask_shape(shape))) + ","
+            shape_str += (np.array2string(int64_array(unmask_shape(shape))) if shape is not None else "Undefined") + ","
             if not is_fully_defined(shape):
                 is_partially_defined = "1"
         message_str = "{fw:" + framework + ",shape:\"" + shape_str[:-1] + "\"}"

model-optimizer/unit_tests/mo/utils/telemetry_utils_test.py

@@ -74,6 +74,18 @@ class TestTelemetryUtils(unittest.TestCase):
         tm.Telemetry.send_event.assert_any_call('mo', 'partially_defined_shape',
                                                 '{partially_defined_shape:1,fw:framework}')
 
+    def test_send_undefined_shapes(self):
+        graph = build_graph({**regular_op('placeholder1', {'shape': None,
+                                                           'type': 'Parameter'}),
+                             **regular_op('mul', {'shape': int64_array([7, 8]), 'type': 'Multiply'})}, [])
+
+        self.init_telemetry_mocks()
+
+        send_shapes_info('framework', graph)
+        tm.Telemetry.send_event.assert_any_call('mo', 'input_shapes', '{fw:framework,shape:"Undefined"}')
+        tm.Telemetry.send_event.assert_any_call('mo', 'partially_defined_shape',
+                                                '{partially_defined_shape:1,fw:framework}')
+
     def test_send_dynamic_shapes_case2(self):
         graph = build_graph({**regular_op('placeholder1', {'shape': int64_array([2, 3, 20, 20]), 'type': 'Parameter'}),
                              **regular_op('placeholder2', {'shape': int64_array([7, 4, 10]), 'type': 'Parameter'}),