From fcd632633b863c190cc8d97a1d1ead02d18d6e7b Mon Sep 17 00:00:00 2001 From: SunAhong1993 Date: Tue, 22 Dec 2020 17:18:33 +0800 Subject: [PATCH] for onnx argmax --- .../op_mapper/dygraph/onnx2paddle/opset9/opset.py | 13 +++++++++++++ .../op_mapper/static/onnx2paddle/opset9/opset.py | 13 +++++++++++++ 2 files changed, 26 insertions(+) diff --git a/x2paddle/op_mapper/dygraph/onnx2paddle/opset9/opset.py b/x2paddle/op_mapper/dygraph/onnx2paddle/opset9/opset.py index f0b51d4..b1c1fd8 100644 --- a/x2paddle/op_mapper/dygraph/onnx2paddle/opset9/opset.py +++ b/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) diff --git a/x2paddle/op_mapper/static/onnx2paddle/opset9/opset.py b/x2paddle/op_mapper/static/onnx2paddle/opset9/opset.py index 53c2870..e4c5439 100644 --- a/x2paddle/op_mapper/static/onnx2paddle/opset9/opset.py +++ b/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 -- GitLab