diff --git a/model-optimizer/extensions/back/ReverseInputChannels.py b/model-optimizer/extensions/back/ReverseInputChannels.py index 50d78fcd95b..e8dd267d460 100644 --- a/model-optimizer/extensions/back/ReverseInputChannels.py +++ b/model-optimizer/extensions/back/ReverseInputChannels.py @@ -310,7 +310,13 @@ class ReverseChannelsPropagationUp(BackReplacementPattern): reverse_channels_copy = reverse_channels.copy_node({'axis': np.array(axis)}) src = port.get_connection().get_source() - port.get_connection().set_source(reverse_channels_copy.out_port(0)) + if src.node.soft_get('type') == 'Parameter': + # For Parameter nodes tensor debug attributes should not move to the last node + # of subgraph. It is needed for the proper mapping of input framework name. + # For this reason "source" mode is used to keep tensor debug attributes at Parameter node. + port.get_connection().set_source(reverse_channels_copy.out_port(0), attributes_save_mode="source") + else: + port.get_connection().set_source(reverse_channels_copy.out_port(0)) src.connect(reverse_channels_copy.in_port(0)) copies.append(reverse_channels_copy) diff --git a/model-optimizer/unit_tests/extensions/back/ReverseInputChannels_test.py b/model-optimizer/unit_tests/extensions/back/ReverseInputChannels_test.py new file mode 100644 index 00000000000..057f84010e3 --- /dev/null +++ b/model-optimizer/unit_tests/extensions/back/ReverseInputChannels_test.py @@ -0,0 +1,49 @@ +# Copyright (C) 2018-2021 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +import unittest + +from extensions.back.ReverseInputChannels import ReverseChannelsPropagationUp +from mo.graph.graph import Node, Graph +from unit_tests.utils.graph import build_graph, result, connect, regular_op_with_shaped_data + +nodes = { + **regular_op_with_shaped_data('placeholder1', [1, 3, 10, 10], {'type': 'Parameter'}), + **regular_op_with_shaped_data('placeholder2', [1, 1, 1, 1], {'type': 'Parameter'}), + + **regular_op_with_shaped_data('mul', [1, 3, 10, 10], {'type': 'Multiply'}), + **regular_op_with_shaped_data('reverse_channels', [1, 3, 10, 10], {'type': 'ReverseChannels', 'axis': 1}), + + **result('result'), +} + + +class ReverseInputChannelsTest(unittest.TestCase): + def check_graph_attrs(self, graph: Graph, parameter_node_names: list): + for node in graph.get_op_nodes(): + if node.soft_get('name') in parameter_node_names: + self.assertTrue(node.soft_get('type') == 'Parameter') + out_node = node.out_node(0) + self.assertTrue(out_node['fw_tensor_debug_info'] == ['fw_name', 0]) + else: + for idx in node.out_nodes(): + out_node = node.out_node(idx) + self.assertFalse('fw_tensor_debug_info' in out_node) + + def set_graph_attrs(self, graph: Graph, parameter_node_names: list): + for node in graph.get_op_nodes(): + if node.soft_get('name') in parameter_node_names: + self.assertTrue(node.soft_get('type') == 'Parameter') + out_node = node.out_node(0) + out_node['fw_tensor_debug_info'] = ['fw_name', 0] + + def test_lift_up_through_eltwise(self): + graph = build_graph(nodes, [*connect('placeholder1', '0:mul'), *connect('placeholder2', '1:mul'), + *connect('mul', 'reverse_channels'), *connect('reverse_channels', 'result')]) + self.set_graph_attrs(graph, ['placeholder1', 'placeholder2']) + + node = Node(graph, 'mul') + reverse_channels = Node(graph, 'reverse_channels') + + ReverseChannelsPropagationUp.lift_up_through_eltwise(node, reverse_channels) + self.check_graph_attrs(graph, ['placeholder1', 'placeholder2'])