Files
openvino/model-optimizer/extensions/back/TopKNormalizer.py
T
Evgeny Lazarev 970b1301b5 Cleanup IR v7 from the MO (#1008)
* Removed back phase transformations related to IRv7

* Fixed setting value for the input port using the 'set_value' method

* Removed front and middle phase transformations related to IRv7

* Cleanup the rest of the Model Optimizer transformations from IRv7 specific transformations

* Final cleanup of the deprecated IR v7 related code

* Removed 'blobs_as_input' usage in the Model Optimizer.

* Removed function '_fuse_add' from the Model Optimizer since it is not used anymore.

* Removed 'keep_in_IR' node attribute for FakeQuantize ops in the MO

* Disabled failing gpu_engine.user_context test
2020-06-22 11:52:00 +03:00

73 lines
2.7 KiB
Python

"""
Copyright (C) 2018-2020 Intel Corporation
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
from extensions.back.Reshape0DToSqueeze import Reshape0DToSqueeze
from extensions.back.ScalarConstNormalize import ScalarNormalize
from mo.back.replacement import BackReplacementPattern
from mo.front.common.partial_infer.utils import int64_array
from mo.front.tf.graph_utils import create_op_node_with_second_input
from mo.graph.graph import Graph, Node
from mo.ops.reshape import Reshape
from mo.ops.result import Result
class TopKNormalizer(BackReplacementPattern):
"""
The transformation converts the second input to the TopK layer from 0D to 1D.
Also the transformation adds the Result Op if there are no consumers of TopK outputs. However the Result for output
with values is not added if the node has attribute 'remove_values_output' which is set to True for Caffe models
where ArgMax does not have separate output with values.
TODO this pass should be removed when IE supports 0D tensors.
"""
enabled = True
def run_after(self):
return [ScalarNormalize]
def run_before(self):
return [Reshape0DToSqueeze]
@staticmethod
def pattern():
return dict(
nodes=[('result', {'type': 'TopK'})],
edges=[],
)
@staticmethod
def replace_pattern(graph: Graph, match: dict):
node = match['result']
reshape = create_op_node_with_second_input(graph, Reshape, int64_array([]), {'override_output_shape': True})
node.in_port(1).get_connection().insert_node(reshape)
TopKNormalizer.normalize_outputs(node)
@staticmethod
def normalize_outputs(node: Node):
"""
This function adds missed outputs for TopK node.
"""
if node.out_port(0).disconnected():
output = Result(node.graph, {'name': node.name + '/Result_port_0/',
'remove_from_xml': node.has_and_set('remove_values_output')}).create_node()
node.out_port(0).get_connection().set_destination(output.in_port(0))
if node.out_port(1).disconnected():
output = Result(node.graph, {'name': node.name + '/Result_port_1/'}).create_node()
node.out_port(1).get_connection().set_destination(output.in_port(0))