Add remove_value_output attribute to onnx ArgMax extractor (#4831)

* add remove_value_output to onnx argmax extractor

* update TopK extender

* update copyright

* change condition in topk extender based on connection logic

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

@@ -35,7 +35,9 @@ class ArgMaxFrontExtractor(FrontExtractorOp):
             'out_max_val' : 0,
 
             # Set attribute to trigger ArgMax replacer in case do not keep the dimension
-            'keepdims': keepdims
+            'keepdims': keepdims,
+
+            'remove_values_output': True
         }
 
         ArgMaxOp.update_node_stat(node, attrs)

model-optimizer/mo/utils/ir_reader/extenders/topk_extender.py

@@ -1,5 +1,5 @@
 """
- Copyright (c) 2020 Intel Corporation
+ Copyright (c) 2020-2021 Intel Corporation
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -24,7 +24,7 @@ class TopKExtender(Extender):
 
     @staticmethod
     def extend(op: Node):
-        if op.graph.graph['cmd_params'].framework in ('tf', 'caffe'):
+        if op.out_port(0).disconnected():
             op['remove_values_output'] = True
         if op.has_valid('index_element_type'):
             op['index_element_type'] = destination_type_to_np_data_type(op.index_element_type)