diff --git a/model-optimizer/mo/graph/port.py b/model-optimizer/mo/graph/port.py index 9a5132e387c..80abb66684a 100644 --- a/model-optimizer/mo/graph/port.py +++ b/model-optimizer/mo/graph/port.py @@ -101,7 +101,7 @@ class Port: else: data_node = self.node.out_node(self.idx, control_flow=self.control_flow) assert data_node.value is None or \ - np.array_equal(data_node.soft_get('force_shape', data_node.shape), int64_array(shape)) + np.array_equal(data_node.soft_get('force_shape', data_node.shape), int64_array(shape)) self.node.out_node(self.idx, control_flow=self.control_flow).shape = int64_array(shape) def _get_value(self): @@ -263,25 +263,37 @@ class Port: return consumer_ports def get_tensor_names(self, port_renumber: bool = False): - def get_tensor_names_list(attrs): - tensor_names_list = [] + """ + Gets sorted tensor names list. + :param port_renumber: defines whether data node index should be calculated considering port renumbering. + """ + tensor_debug_info = self.get_tensor_debug_info(port_renumber) + tensor_names_list = [] + for attr in tensor_debug_info: + if attr is not None and len(attr) >= 2: + tensor_name = attr[1] + if tensor_name is not None and len(tensor_name) > 0: + tensor_names_list.append(tensor_name.replace(',', '\\,')) + return sorted(tensor_names_list) + + def get_tensor_debug_info(self, port_renumber: bool = False): + """ + Gets tensor debug info attribute. + :param port_renumber: defines whether data node index should be calculated considering port renumbering. + """ + def get_tensor_debug_info_from_attrs(attrs): if 'fw_tensor_debug_info' in attrs: - if attrs['fw_tensor_debug_info'] is None: - return tensor_names_list - for attr in attrs['fw_tensor_debug_info']: - if attr is not None and len(attr) >= 2: - tensor_name = attr[1] - if tensor_name is not None and len(tensor_name) > 0: - tensor_names_list.append(tensor_name.replace(',', '\\,')) - return tensor_names_list + if attrs['fw_tensor_debug_info'] is not None: + return attrs['fw_tensor_debug_info'] + return [] - assert self.type != 'in', "Can't get tensor names for input port at {} node".format(self.node.name) + assert self.type != 'in', "Can't get tensor debug info for input port at {} node".format(self.node.name) - fw_names = [] + fw_debug_info = [] if self.node.graph.stage == 'front': if self.idx in self.node.out_edges(): out_edge = self.node.out_edge(self.idx) - fw_names += get_tensor_names_list(out_edge) + fw_debug_info += get_tensor_debug_info_from_attrs(out_edge) else: # before port renumbering we use sequential numbering node_idx = self.idx @@ -293,8 +305,9 @@ class Port: if node_idx in self.node.out_nodes(): out_node = self.node.out_node(node_idx) - fw_names += get_tensor_names_list(out_node.attrs()) - return sorted(fw_names) + fw_debug_info += get_tensor_debug_info_from_attrs(out_node.attrs()) + return fw_debug_info + def disconnect(self): if self.type == 'out': diff --git a/model-optimizer/mo/ops/op.py b/model-optimizer/mo/ops/op.py index 6f45fa76e6f..7af2448ea0e 100644 --- a/model-optimizer/mo/ops/op.py +++ b/model-optimizer/mo/ops/op.py @@ -121,13 +121,20 @@ class Op(object): if attrs is None: attrs = dict() new_node = self.add_node(attrs) - # Missed careful handling of debug information for i, inp in enumerate(inputs): edge_attr = {'in': i, 'out': inp[1], 'in_attrs': ['in', 'permutation'], 'out_attrs': ['out', 'permutation'], 'data_attrs': []} if not inp[0].has_valid('kind') or inp[0].kind == 'op' \ else {'in': i, 'in_attrs': ['in', 'permutation']} + + # handling of debug information + if inp[0].has_port('out', inp[1]): + debug_info = inp[0].out_port(inp[1]).get_tensor_debug_info() + if debug_info is not None and len(debug_info) > 0: + edge_attr.update({'fw_tensor_debug_info': debug_info}) + edge_attr['data_attrs'].append('fw_tensor_debug_info') + if edge_attrs is not None: edge_attr.update(edge_attrs) new_node.add_input_port(i, skip_if_exist=True) diff --git a/model-optimizer/unit_tests/mo/ops/op_test.py b/model-optimizer/unit_tests/mo/ops/op_test.py new file mode 100644 index 00000000000..dbdde9b4cf8 --- /dev/null +++ b/model-optimizer/unit_tests/mo/ops/op_test.py @@ -0,0 +1,30 @@ +# Copyright (C) 2018-2021 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +import unittest + +from extensions.ops.lstm_cell import LSTMCell +from mo.graph.graph import Node +from unit_tests.utils.graph import build_graph, regular_op + +nodes = { + **regular_op('Op1', {'type': 'Op1', 'kind': 'op', 'op': 'Op1'}), + **regular_op('Op2', {'type': 'Op2', 'kind': 'op', 'op': 'Op2'}), + **regular_op('Op3', {'type': 'Op3', 'kind': 'op', 'op': 'Op3'}), +} + + +class TestOp(unittest.TestCase): + def test_create_node(self): + graph = build_graph(nodes, [('Op1', 'Op3', {'in': 0, 'out': 0, 'fw_tensor_debug_info': [('Op1', 'Op1')]}), + ('Op2', 'Op3', {'in': 1, 'out': 0, 'fw_tensor_debug_info': [('Op2', 'Op2')]})]) + graph.stage = 'front' + input1 = Node(graph, 'Op1') + input2 = Node(graph, 'Op2') + inputs = [(input1, 0), (input2, 0)] + + lstm_op = LSTMCell(graph, dict(name='LSTMCell')) + _ = lstm_op.create_node(inputs) + + self.assertTrue(input1.out_edge(0)['fw_tensor_debug_info'] == [('Op1', 'Op1')]) + self.assertTrue(input2.out_edge(0)['fw_tensor_debug_info'] == [('Op2', 'Op2')])