Implementation of ArgMin ONNX + TF extractors (#5126)

* implement argmin extractors

* reconsidering argmax to topk

* arg ops refactoring

* rename ArgMaxToTopK

* added unittests

* update docs

* move unittest file to new folder

* conversations resolving

* revert changes with argmax.py, move argmin op to a new file

* rename ArgMaxSqueeze

* updated BOM file

* little fix

* code refactoring in ArgMaxOp, updated unittests

Co-authored-by: yegor.kruglov <ykruglov@nnlvdp-mkaglins.inn.intel.com>

model-optimizer/extensions/front/ArgOpsSqueeze.py

@@ -7,22 +7,21 @@ from mo.ops.const import Const
 from mo.ops.squeeze import Squeeze
 
 
-class ArgMaxSqueeze(FrontReplacementSubgraph):
+class ArgOpsSqueeze(FrontReplacementSubgraph):
     """
-        In some frameworks ArgMax operation has keepdims attribute that indicates whether to stay a dimension along
-        which maximum is computed or not. In case of keepdims=0 this dimension should be removed but ArgMax operation in
-        IR format is not designed to cover this case. So we should additionally add Squeeze operation right after ArgMax
-        for this case.
+        In some frameworks ArgMax/ArgMin operation has keepdims attribute that indicates whether to stay a dimension
+        along which maximum is computed or not. In case of keepdims=0 this dimension should be removed but ArgMax/ArgMin
+        operation in IR format is not designed to cover this case. So we should additionally add Squeeze operation right
+        after ArgMax/ArgMin for this case.
     """
-    op = "ArgMax"
     enabled = True
 
     def pattern(self):
-        return dict(nodes=[('argmax', dict(op='ArgMax', keepdims=0))],
+        return dict(nodes=[('node', dict(op=lambda x: x in ['ArgMax', 'ArgMin'], keepdims=0))],
                     edges=[])
 
     def replace_sub_graph(self, graph: Graph, match: dict):
-        node = match['argmax']
+        node = match['node']
 
         connected_ports = [port for port in node.in_ports().values() if not port.disconnected()]
         squeeze_node = Squeeze(graph, dict()).create_node([], dict(name=node.name + '/Squeeze'))

model-optimizer/extensions/front/onnx/argmin_ext.py

@@ -0,0 +1,26 @@
+# Copyright (C) 2018-2021 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from extensions.ops.argmin import ArgMinOp
+from mo.front.extractor import FrontExtractorOp
+from mo.front.onnx.extractors.utils import onnx_attr
+
+
+class ArgMinFrontExtractor(FrontExtractorOp):
+    op = 'ArgMin'
+    enabled = True
+
+    @classmethod
+    def extract(cls, node):
+        keepdims = onnx_attr(node, 'keepdims', 'i', default=1)
+        axis = onnx_attr(node, 'axis', 'i', default=0)
+
+        attrs = {
+            'axis': axis,
+            'top_k': 1,
+            'keepdims': keepdims,
+            'remove_values_output': True
+        }
+
+        ArgMinOp.update_node_stat(node, attrs)
+        return cls.enabled

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

@@ -16,6 +16,6 @@ class ArgMaxFrontExtractor(FrontExtractorOp):
     def extract(cls, node):
         ArgMaxOp.update_node_stat(node, {'out_max_val': 0, 'top_k': 1, 'axis': None,
                                          'dim_attrs': ['axis'], 'keepdims': 0, 'remove_values_output': True,
-                                         'output_type': tf_dtype_extractor(node.pb.attr['out_type'].type, np.int64),
+                                         'output_type': tf_dtype_extractor(node.pb.attr['output_type'].type, np.int64),
                                          })
         return cls.enabled

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

@@ -0,0 +1,25 @@
+# Copyright (C) 2018-2021 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+
+from extensions.ops.argmin import ArgMinOp
+from mo.front.extractor import FrontExtractorOp
+from mo.front.tf.extractors.utils import tf_dtype_extractor
+
+
+class ArgMinFrontExtractor(FrontExtractorOp):
+    op = 'ArgMin'
+    enabled = True
+
+    @classmethod
+    def extract(cls, node):
+        attrs = {
+            'top_k': 1,
+            'axis': None,
+            'keepdims': 0,
+            'remove_values_output': True,
+            'output_type': tf_dtype_extractor(node.pb.attr['output_type'].type, np.int64)
+        }
+        ArgMinOp.update_node_stat(node, attrs)
+        return cls.enabled

model-optimizer/extensions/middle/ArgOpsToTopK.py

@@ -8,24 +8,24 @@ from mo.ops.concat import Concat
 from mo.ops.const import Const
 
 
-class ArgMaxToTopK(MiddleReplacementPattern):
+class ArgOpsToTopK(MiddleReplacementPattern):
     """
-        The transformation replaces ArgMax with the TopK layer.
+        The transformation replaces ArgMax/ArgMin with the TopK layer.
     """
-    op = "ArgMax"
+
     enabled = True
     force_clean_up = True
 
     def pattern(self):
         return dict(
             nodes=[
-                ('argmax', dict(op='ArgMax')),
+                ('node', dict(op=lambda x: x in ['ArgMax', 'ArgMin'])),
             ],
             edges=[]
         )
 
     def replace_pattern(self, graph: Graph, match: dict):
-        node = match['argmax']
+        node = match['node']
         node_name = node.soft_get('name', node.id)
 
         connected_ports = [port for port in node.in_ports().values() if not port.disconnected()]
@@ -36,7 +36,9 @@ class ArgMaxToTopK(MiddleReplacementPattern):
 
         assert axis is not None, 'The "axis" should be defined for node "{}"'.format(node_name)
         assert node.has_and_set('output_type'), 'The data type is not set for node "{}"'.format(node_name)
-        topk_node = TopK(graph, {'axis': axis, 'mode': 'max', 'sort': 'index',
+
+        topk_mode = 'max' if node.op == 'ArgMax' else 'min'
+        topk_node = TopK(graph, {'axis': axis, 'mode': topk_mode, 'sort': 'index',
                                  'remove_values_output': node.has_and_set('remove_values_output'),
                                  'index_element_type': node.output_type}).create_node()
         node.in_port(0).get_connection().set_destination(topk_node.in_port(0))
@@ -47,7 +49,7 @@ class ArgMaxToTopK(MiddleReplacementPattern):
             topk_node.out_port(0).connect(concat_node.in_port(1))  # indices
             topk_node.out_port(1).connect(concat_node.in_port(0))  # values
             if not node.out_port(0).disconnected():
-                node.out_port(0).get_connection().set_source(concat_node.out_port(1))
+                node.out_port(0).get_connection().set_source(concat_node.out_port(0))
         else:
             if not node.out_port(0).disconnected():
                 node.out_port(0).get_connection().set_source(topk_node.out_port(1))

model-optimizer/extensions/ops/argmax.py

@@ -6,18 +6,57 @@ import logging as log
 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
 from mo.graph.graph import Node, Graph
 from mo.ops.op import Op, PermuteAttrs
 
 
+def arg_ops_infer(node: Node):
+    shape = node.in_port(0).data.get_shape()
+    node_name = node.soft_get('name', node.id)
+    assert shape is not None, "Input shape for the node {} is None".format(node_name)
+
+    # there are two inputs in TensorFlow. The second input is the axis for ArgMax
+    connected_in_ports = [port for port in node.in_ports().values() if not port.disconnected()]
+    if len(connected_in_ports) == 2:
+        axis = node.in_port(1).data.get_value()
+        if axis is None:
+            log.debug('The second argument to {} is None'.format(node.soft_get('name', node.id)))
+            return
+        node.axis = axis
+        # remove the unnecessary input
+        node.in_port(1).disconnect()
+
+    num_top_axes = shape.size
+    if num_top_axes < 3:
+        num_top_axes = 3
+
+    out_shape = np.ones(num_top_axes, dtype=np.int64)
+
+    if node.has_valid('axis'):
+        axis = get_canonical_axis_index(shape, node.axis)
+        node.axis = axis
+        out_shape = int64_array(shape)
+        out_shape[axis] = node.top_k
+        PermuteAttrs.create_permute_attrs(node, attrs=[('axis', 'input:0')])
+    else:
+        out_shape[0] = shape[0]
+        out_shape[2] = node.top_k
+        if node.has_and_set('out_max_val'):
+            out_shape[1] = 2
+
+    node.out_port(0).data.set_shape(out_shape)
+
+
 class ArgMaxOp(Op):
     op = 'ArgMax'
+    enabled = False
 
     def __init__(self, graph: Graph, attrs: dict):
         mandatory_props = {
-            'type': __class__.op,
-            'op': __class__.op,
-            'infer': ArgMaxOp.argmax_infer,
+            'type': None,
+            'op': self.op,
+            'infer': arg_ops_infer,
             'output_type': np.int64,
             'in_ports_count': 2,
             'out_ports_count': 1,
@@ -30,38 +69,3 @@ class ArgMaxOp(Op):
             'top_k',
             'axis',
         ]
-
-    @staticmethod
-    def argmax_infer(node: Node):
-        shape = node.in_node(0).shape
-        if shape is None:
-            return
-
-        # there are two inputs in TensorFlow. The second input is the axis for ArgMax
-        if len(node.in_nodes()) == 2:
-            if node.in_node(1).value is None:
-                log.debug('The second argument to ArgMax is None')
-                return
-            node.axis = node.in_node(1).value.item()
-            # remove the unnecessary input
-            node.graph.remove_edge(node.in_node(1).id, node.id)
-
-        num_top_axes = shape.size
-        if num_top_axes < 3:
-            num_top_axes = 3
-
-        out_shape = np.ones(num_top_axes, dtype=int)
-
-        if node.has_valid('axis'):
-            axis = get_canonical_axis_index(shape, node.axis)
-            node.axis = axis
-            out_shape = np.array(shape)
-            out_shape[axis] = node.top_k
-            PermuteAttrs.create_permute_attrs(node, attrs=[('axis', 'input:0')])
-        else:
-            out_shape[0] = shape[0]
-            out_shape[2] = node.top_k
-            if node.out_max_val:
-                out_shape[1] = 2
-
-        node.out_node().shape = out_shape

model-optimizer/extensions/ops/argmin.py

@@ -0,0 +1,30 @@
+# Copyright (C) 2018-2021 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+
+from extensions.ops.argmax import arg_ops_infer
+from mo.graph.graph import Graph
+from mo.ops.op import Op
+
+
+class ArgMinOp(Op):
+    op = 'ArgMin'
+    enabled = False
+
+    def __init__(self, graph: Graph, attrs: dict):
+        mandatory_props = {
+            'type': None,
+            'op': self.op,
+            'infer': arg_ops_infer,
+            'output_type': np.int64,
+            'in_ports_count': 2,
+            'out_ports_count': 1,
+        }
+        super().__init__(graph, mandatory_props, attrs)
+
+    def supported_attrs(self):
+        return [
+            'top_k',
+            'axis',
+        ]