Nigel
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

!37370 [assistant][ops][I4CRJB/I4CRJC] Add ArgmaxV2 and ArgminV3 

Merge pull request !37370 from chenchen/argmax
This commit is contained in:
i-robot 2022-09-23 03:57:45 +00:00 committed by Gitee
parent de3282e2e4 5f6135f290
commit 73e4469181
No known key found for this signature in database
GPG Key ID: 173E9B9CA92EEF8F
9 changed files with 196 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mindspore/ccsrc/plugin/device/ascend/kernel/aicpu/aicpu_attr_to_input_registry.cc
View File

@ -43,6 +43,8 @@ std::map<string, std::vector<std::pair<string, size_t>>> AicpuOpAttrToInputMap =
{prim::kPrimReduceProd->name(), {{"axis", 1}}},
{prim::kPrimReverseV2->name(), {{"axis", 1}}},
{prim::kPrimBroadcastTo->name(), {{"shape", 1}}},
{prim::kPrimArgMax->name(), {{"axis", 1}}},
{prim::kPrimArgMin->name(), {{"axis", 1}}},
{prim::kPrimUnsortedSegmentProd->name(), {{"num_segments", 2}}},
{prim::kPrimUnsortedSegmentSum->name(), {{"num_segments", 2}}}};

6
mindspore/ccsrc/plugin/device/ascend/kernel/aicpu/aicpu_util.h
View File

@ -171,6 +171,8 @@ constexpr auto kZerosLike = "ZerosLike";
constexpr auto kEqual = "Equal";
constexpr auto kOnesLike = "OnesLike";
constexpr auto kSign = "Sign";
constexpr auto kArgmax = "Argmax";
constexpr auto kArgmin = "Argmin";
const std::set<std::string> kCpuKernelOps{kIdentity,
kMaskedSelect,
@ -295,7 +297,9 @@ const std::map<std::string, std::string> kOpNameToAicpuOpNameMap{
{kTensorScatterElements, "ScatterElements"},
{kACos, "Acos"},
{kHSigmoid, "HardSigmoid"},
{kHSigmoidGrad, "HardSigmoidGrad"}};
{kHSigmoidGrad, "HardSigmoidGrad"},
{kArgmax, "ArgMax"},
{kArgmin, "ArgMin"}};
struct AicpuParamHead {
uint32_t length; // Total length: include cunstom message
uint32_t ioAddrNum; // Input and output address number

4
mindspore/core/abstract/ops/primitive_infer_map.cc
View File

@ -126,7 +126,9 @@ PrimShapeDependMap &GetHostDependsMap() {
static const auto &kTraceGrad = prim::kPrimTraceGrad->name();
static const auto &kSetSize = prim::kPrimSetSize->name();
// Common host depends.
static PrimShapeDependMap host_depends{{kExtractGlimpse, ShapeSet{1}},
static PrimShapeDependMap host_depends{{prim::kPrimArgMax->name(), ShapeSet{1}},
{prim::kPrimArgMin->name(), ShapeSet{1}},
{kExtractGlimpse, ShapeSet{1}},
{kSegmentMax, ShapeSet{1}},
{kSegmentMin, ShapeSet{1}},
{kSegmentSum, ShapeSet{1}},

2
mindspore/core/ops/arg_max.cc
View File

@ -15,6 +15,8 @@
*/
#include "ops/arg_max.h"
#include <set>
#include "mindapi/ir/type.h"
#include "utils/check_convert_utils.h"
#include "ops/op_utils.h"

3
mindspore/core/ops/arg_max.h
View File

@ -1,5 +1,5 @@
/**
* Copyright 2020 Huawei Technologies Co., Ltd
* Copyright 2022 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
@ -53,6 +53,7 @@ class MIND_API Argmax : public BaseOperator {
};
abstract::AbstractBasePtr ArgMaxInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive,
const std::vector<abstract::AbstractBasePtr> &input_args);
using PrimArgMax = std::shared_ptr<Argmax>;
} // namespace ops
} // namespace mindspore

3
mindspore/core/ops/arg_min.h
View File

@ -1,5 +1,5 @@
/**
* Copyright 2020 Huawei Technologies Co., Ltd
* Copyright 2022 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
@ -53,6 +53,7 @@ class MIND_API ArgMin : public BaseOperator {
};
abstract::AbstractBasePtr ArgMinInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive,
const std::vector<abstract::AbstractBasePtr> &input_args);
using PrimArgMin = std::shared_ptr<ArgMin>;
} // namespace ops
} // namespace mindspore

75
mindspore/python/mindspore/ops/_op_impl/aicpu/arg_max.py Normal file
View File

@ -0,0 +1,75 @@
# Copyright 2022 Huawei Technologies Co., Ltd
#
# 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.
# ============================================================================
"""Argmax op"""
from mindspore.ops.op_info_register import op_info_register, AiCPURegOp, DataType
arg_max_info = AiCPURegOp("Argmax") \
.fusion_type("OPAQUE") \
.attr("cust_aicpu", "str") \
.input(0, "x", "required") \
.input(1, "dimension", "required") \
.output(0, "y", "required") \
.dtype_format(DataType.F16_Default, DataType.I32_Default, DataType.I32_Default) \
.dtype_format(DataType.F16_Default, DataType.I32_Default, DataType.I64_Default) \
.dtype_format(DataType.F16_Default, DataType.I64_Default, DataType.I32_Default) \
.dtype_format(DataType.F16_Default, DataType.I64_Default, DataType.I64_Default) \
.dtype_format(DataType.F32_Default, DataType.I32_Default, DataType.I32_Default) \
.dtype_format(DataType.F32_Default, DataType.I32_Default, DataType.I64_Default) \
.dtype_format(DataType.F32_Default, DataType.I64_Default, DataType.I32_Default) \
.dtype_format(DataType.F32_Default, DataType.I64_Default, DataType.I64_Default) \
.dtype_format(DataType.F64_Default, DataType.I32_Default, DataType.I32_Default) \
.dtype_format(DataType.F64_Default, DataType.I32_Default, DataType.I64_Default) \
.dtype_format(DataType.F64_Default, DataType.I64_Default, DataType.I32_Default) \
.dtype_format(DataType.F64_Default, DataType.I64_Default, DataType.I64_Default) \
.dtype_format(DataType.I8_Default, DataType.I32_Default, DataType.I32_Default) \
.dtype_format(DataType.I8_Default, DataType.I32_Default, DataType.I64_Default) \
.dtype_format(DataType.I8_Default, DataType.I64_Default, DataType.I32_Default) \
.dtype_format(DataType.I8_Default, DataType.I64_Default, DataType.I64_Default) \
.dtype_format(DataType.I16_Default, DataType.I32_Default, DataType.I32_Default) \
.dtype_format(DataType.I16_Default, DataType.I32_Default, DataType.I64_Default) \
.dtype_format(DataType.I16_Default, DataType.I64_Default, DataType.I32_Default) \
.dtype_format(DataType.I16_Default, DataType.I64_Default, DataType.I64_Default) \
.dtype_format(DataType.I32_Default, DataType.I32_Default, DataType.I32_Default) \
.dtype_format(DataType.I32_Default, DataType.I32_Default, DataType.I64_Default) \
.dtype_format(DataType.I32_Default, DataType.I64_Default, DataType.I32_Default) \
.dtype_format(DataType.I32_Default, DataType.I64_Default, DataType.I64_Default) \
.dtype_format(DataType.I64_Default, DataType.I32_Default, DataType.I32_Default) \
.dtype_format(DataType.I64_Default, DataType.I32_Default, DataType.I64_Default) \
.dtype_format(DataType.I64_Default, DataType.I64_Default, DataType.I32_Default) \
.dtype_format(DataType.I64_Default, DataType.I64_Default, DataType.I64_Default) \
.dtype_format(DataType.U8_Default, DataType.I32_Default, DataType.I32_Default) \
.dtype_format(DataType.U8_Default, DataType.I32_Default, DataType.I64_Default) \
.dtype_format(DataType.U8_Default, DataType.I64_Default, DataType.I32_Default) \
.dtype_format(DataType.U8_Default, DataType.I64_Default, DataType.I64_Default) \
.dtype_format(DataType.U16_Default, DataType.I32_Default, DataType.I32_Default) \
.dtype_format(DataType.U16_Default, DataType.I32_Default, DataType.I64_Default) \
.dtype_format(DataType.U16_Default, DataType.I64_Default, DataType.I32_Default) \
.dtype_format(DataType.U16_Default, DataType.I64_Default, DataType.I64_Default) \
.dtype_format(DataType.U32_Default, DataType.I32_Default, DataType.I32_Default) \
.dtype_format(DataType.U32_Default, DataType.I32_Default, DataType.I64_Default) \
.dtype_format(DataType.U32_Default, DataType.I64_Default, DataType.I32_Default) \
.dtype_format(DataType.U32_Default, DataType.I64_Default, DataType.I64_Default) \
.dtype_format(DataType.U64_Default, DataType.I32_Default, DataType.I32_Default) \
.dtype_format(DataType.U64_Default, DataType.I32_Default, DataType.I64_Default) \
.dtype_format(DataType.U64_Default, DataType.I64_Default, DataType.I32_Default) \
.dtype_format(DataType.U64_Default, DataType.I64_Default, DataType.I64_Default) \
.get_op_info()
@op_info_register(arg_max_info)
def _arg_max_aicpu():
"""Argmax aicpu register"""
return

75
mindspore/python/mindspore/ops/_op_impl/aicpu/arg_min.py Normal file
View File

@ -0,0 +1,75 @@
# Copyright 2022 Huawei Technologies Co., Ltd
#
# 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.
# ============================================================================
"""Argmin op"""
from mindspore.ops.op_info_register import op_info_register, AiCPURegOp, DataType
arg_min_info = AiCPURegOp("Argmin") \
.fusion_type("OPAQUE") \
.attr("cust_aicpu", "str") \
.input(0, "x", "required") \
.input(1, "dimension", "required") \
.output(0, "y", "required") \
.dtype_format(DataType.F16_Default, DataType.I32_Default, DataType.I32_Default) \
.dtype_format(DataType.F16_Default, DataType.I32_Default, DataType.I64_Default) \
.dtype_format(DataType.F16_Default, DataType.I64_Default, DataType.I32_Default) \
.dtype_format(DataType.F16_Default, DataType.I64_Default, DataType.I64_Default) \
.dtype_format(DataType.F32_Default, DataType.I32_Default, DataType.I32_Default) \
.dtype_format(DataType.F32_Default, DataType.I32_Default, DataType.I64_Default) \
.dtype_format(DataType.F32_Default, DataType.I64_Default, DataType.I32_Default) \
.dtype_format(DataType.F32_Default, DataType.I64_Default, DataType.I64_Default) \
.dtype_format(DataType.F64_Default, DataType.I32_Default, DataType.I32_Default) \
.dtype_format(DataType.F64_Default, DataType.I32_Default, DataType.I64_Default) \
.dtype_format(DataType.F64_Default, DataType.I64_Default, DataType.I32_Default) \
.dtype_format(DataType.F64_Default, DataType.I64_Default, DataType.I64_Default) \
.dtype_format(DataType.I8_Default, DataType.I32_Default, DataType.I32_Default) \
.dtype_format(DataType.I8_Default, DataType.I32_Default, DataType.I64_Default) \
.dtype_format(DataType.I8_Default, DataType.I64_Default, DataType.I32_Default) \
.dtype_format(DataType.I8_Default, DataType.I64_Default, DataType.I64_Default) \
.dtype_format(DataType.I16_Default, DataType.I32_Default, DataType.I32_Default) \
.dtype_format(DataType.I16_Default, DataType.I32_Default, DataType.I64_Default) \
.dtype_format(DataType.I16_Default, DataType.I64_Default, DataType.I32_Default) \
.dtype_format(DataType.I16_Default, DataType.I64_Default, DataType.I64_Default) \
.dtype_format(DataType.I32_Default, DataType.I32_Default, DataType.I32_Default) \
.dtype_format(DataType.I32_Default, DataType.I32_Default, DataType.I64_Default) \
.dtype_format(DataType.I32_Default, DataType.I64_Default, DataType.I32_Default) \
.dtype_format(DataType.I32_Default, DataType.I64_Default, DataType.I64_Default) \
.dtype_format(DataType.I64_Default, DataType.I32_Default, DataType.I32_Default) \
.dtype_format(DataType.I64_Default, DataType.I32_Default, DataType.I64_Default) \
.dtype_format(DataType.I64_Default, DataType.I64_Default, DataType.I32_Default) \
.dtype_format(DataType.I64_Default, DataType.I64_Default, DataType.I64_Default) \
.dtype_format(DataType.U8_Default, DataType.I32_Default, DataType.I32_Default) \
.dtype_format(DataType.U8_Default, DataType.I32_Default, DataType.I64_Default) \
.dtype_format(DataType.U8_Default, DataType.I64_Default, DataType.I32_Default) \
.dtype_format(DataType.U8_Default, DataType.I64_Default, DataType.I64_Default) \
.dtype_format(DataType.U16_Default, DataType.I32_Default, DataType.I32_Default) \
.dtype_format(DataType.U16_Default, DataType.I32_Default, DataType.I64_Default) \
.dtype_format(DataType.U16_Default, DataType.I64_Default, DataType.I32_Default) \
.dtype_format(DataType.U16_Default, DataType.I64_Default, DataType.I64_Default) \
.dtype_format(DataType.U32_Default, DataType.I32_Default, DataType.I32_Default) \
.dtype_format(DataType.U32_Default, DataType.I32_Default, DataType.I64_Default) \
.dtype_format(DataType.U32_Default, DataType.I64_Default, DataType.I32_Default) \
.dtype_format(DataType.U32_Default, DataType.I64_Default, DataType.I64_Default) \
.dtype_format(DataType.U64_Default, DataType.I32_Default, DataType.I32_Default) \
.dtype_format(DataType.U64_Default, DataType.I32_Default, DataType.I64_Default) \
.dtype_format(DataType.U64_Default, DataType.I64_Default, DataType.I32_Default) \
.dtype_format(DataType.U64_Default, DataType.I64_Default, DataType.I64_Default) \
.get_op_info()
@op_info_register(arg_min_info)
def _arg_min_aicpu():
"""Argmin aicpu register"""
return

50
mindspore/python/mindspore/ops/operations/array_ops.py
View File

@ -1999,12 +1999,35 @@ class InvertPermutation(PrimitiveWithInfer):
'value': tuple(y)}
class Argmax(PrimitiveWithInfer):
class Argmax(Primitive):
"""
Returns the indices of the maximum value of a tensor across the axis.
Refer to :func:`mindspore.ops.argmax` for more detail.
If the shape of input tensor is :math:`(x_1, ..., x_N)`, the shape of the output tensor will be
:math:`(x_1, ..., x_{axis-1}, x_{axis+1}, ..., x_N)`.
Args:
axis (int): Axis where the Argmax operation applies to. Default: -1.
output_type (:class:`mindspore.dtype`): An optional data type of `mindspore.dtype.int32` and
`mindspore.dtype.int64`. Default: `mindspore.dtype.int32`.
Inputs:
- **input_x** (Tensor) - Input tensor. :math:`(N,*)` where :math:`*` means, any number of additional dimensions.
Support data type list as follows:
- Ascend: Float16, Float32, Float64, Int8, Int16, Int32, Int64, UInt8, UInt16, UInt32, UInt64.
- GPU: Float16, Float32.
- CPU: Float16, Float32, Float64.
Outputs:
Tensor, whose dtype is determined by `output_type`.
Raises:
TypeError: If `axis` is not an int.
TypeError: If `output_type` is neither int32 nor int64.
Supported Platforms:
``Ascend`` ``GPU`` ``CPU``
@ -2020,25 +2043,10 @@ class Argmax(PrimitiveWithInfer):
"""Initialize Argmax"""
self.init_prim_io_names(inputs=['x'], outputs=['output'])
validator.check_value_type("axis", axis, [int], self.name)
validator.check_types_same_and_valid({'output': output_type}, [mstype.int32], self.name)
validator.check_types_same_and_valid({'output': output_type}, [mstype.int32, mstype.int64], self.name)
self.axis = axis
self.add_prim_attr('output_type', output_type)
def infer_shape(self, x_shape):
axis = self.axis
if axis is None:
axis = 0
x_rank = len(x_shape)
validator.check_int_range(axis, -x_rank, x_rank, Rel.INC_LEFT, "axis", self.name)
axis = axis + x_rank if axis < 0 else axis
ouput_shape = [x_shape[i] for i in range(x_rank) if i != axis]
return ouput_shape
def infer_dtype(self, x_dtype):
validator.check_tensor_dtype_valid("input_x", x_dtype, [mstype.float16, mstype.float32, mstype.float64],
self.name)
return mstype.tensor_type(self.output_type)
class Argmin(Primitive):
"""
@ -2049,15 +2057,17 @@ class Argmin(Primitive):
Args:
axis (int): Axis where the Argmin operation applies to. Default: -1.
output_type (:class:`mindspore.dtype`): An optional data type of `mindspore.dtype.int32`.
Default: `mindspore.dtype.int32`.
output_type (:class:`mindspore.dtype`): An optional data type of `mindspore.dtype.int32` and
`mindspore.dtype.int64`. Default: `mindspore.dtype.int32`.
Inputs:
- **input_x** (Tensor) - Input tensor.
The shape is :math:`(N,*)` where :math:`*` means, any number of additional dimensions.
- Ascend: Float16, Float32, Float64, Int8, Int16, Int32, Int64, UInt8, UInt16, UInt32, UInt64.
Outputs:
Tensor, indices of the min value of input tensor across the axis.
Tensor, whose dtype is determined by `output_type`.
Raises:
TypeError: If `axis` is not an int.