Merge pull request #466 from SunAhong1993/onnx

for onnx argmax

x2paddle/op_mapper/dygraph/onnx2paddle/opset9/opset.py

@@ -1637,3 +1637,16 @@ class OpSet9():
             inputs=inputs_dict,
             outputs=[node.name],
             **layer_attrs)
+        
+    @print_mapping_info
+    def ArgMax(self, node):
+        val_x = self.graph.get_input_node(node, idx=0, copy=True)
+        axis = node.get_attr('axis')
+        keepdims = False if node.get_attr('keepdims') == 0 else True
+        layer_attrs = {'axis': axis,
+                      'keepdim': keepdims}
+        self.paddle_graph.add_layer(
+            'paddle.argmax', 
+            inputs={"x": val_x.name}, 
+            outputs=[node.name],
+            **layer_attrs)

x2paddle/op_mapper/static/onnx2paddle/opset9/opset.py

@@ -1576,4 +1576,17 @@ class OpSet9():
             kernel=paddle_op,
             inputs=layer_inputs,
             outputs=[node.name],
+            **layer_attrs)
+        
+    @print_mapping_info
+    def ArgMax(self, node):
+        val_x = self.graph.get_input_node(node, idx=0, copy=True)
+        axis = node.get_attr('axis')
+        keepdims = False if node.get_attr('keepdims') == 0 else True
+        layer_attrs = {'axis': axis,
+                      'keepdim': keepdims}
+        self.paddle_graph.add_layer(
+            'paddle.argmax', 
+            inputs={"x": val_x.name}, 
+            outputs=[node.name],
             **layer_attrs)
\ No newline at end of file