From d76ed25e96ffe503bc055a0c39012bdb4b9520c2 Mon Sep 17 00:00:00 2001 From: hedongdong Date: Wed, 18 Aug 2021 10:15:17 +0800 Subject: [PATCH] [feat][assistant][I3T92I] add new nn operator Roll --- mindspore/core/base/core_ops.h | 2 + mindspore/core/ops/roll.cc | 56 ++++++++++++ mindspore/core/ops/roll.h | 41 +++++++++ mindspore/nn/layer/basic.py | 87 ++++++++++++++++++- .../ops/_grad_experimental/grad_inner_ops.py | 14 +++ mindspore/ops/_op_impl/tbe/__init__.py | 1 + mindspore/ops/_op_impl/tbe/roll.py | 42 +++++++++ mindspore/ops/operations/_inner_ops.py | 66 ++++++++++++++ tests/ut/python/ops/test_ops.py | 4 + 9 files changed, 312 insertions(+), 1 deletion(-) create mode 100644 mindspore/core/ops/roll.cc create mode 100644 mindspore/core/ops/roll.h create mode 100644 mindspore/ops/_op_impl/tbe/roll.py diff --git a/mindspore/core/base/core_ops.h b/mindspore/core/base/core_ops.h index 15f889d64eb..a3098d093d6 100644 --- a/mindspore/core/base/core_ops.h +++ b/mindspore/core/base/core_ops.h @@ -93,6 +93,7 @@ constexpr auto kDropoutDoMask = "DropoutDoMask"; constexpr auto kDropout = "Dropout"; constexpr auto kDropoutGrad = "DropoutGrad"; constexpr auto kConv2DTranspose = "Conv2DTranspose"; +constexpr auto kRoll = "Roll"; // Here list all primitives used in backend or some special primitives used by core. // GetNext @@ -283,6 +284,7 @@ inline const PrimitivePtr kPrimCTCLossV2Grad = std::make_shared("CTCL inline const PrimitivePtr kPrimCTCLoss = std::make_shared(kCTCLoss); inline const PrimitivePtr kPrimFullConnection = std::make_shared("FullConnection"); inline const PrimitivePtr kPrimConv2DTranspose = std::make_shared(kConv2DTranspose); +inline const PrimitivePtr kPrimRoll = std::make_shared(kRoll); inline const PrimitivePtr kPrimGroupConv2DGradInput = std::make_shared("GroupConv2DGradInput"); inline const PrimitivePtr kPrimBatchNorm = std::make_shared("BatchNorm"); inline const PrimitivePtr kPrimBatchNormGrad = std::make_shared("BatchNormGrad"); diff --git a/mindspore/core/ops/roll.cc b/mindspore/core/ops/roll.cc new file mode 100644 index 00000000000..41fbd33b276 --- /dev/null +++ b/mindspore/core/ops/roll.cc @@ -0,0 +1,56 @@ +/** + * Copyright 2021 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. + */ + +#include "ops/roll.h" +#include + +#include "ops/op_utils.h" +#include "utils/check_convert_utils.h" +#include "utils/tensor_construct_utils.h" +#include "abstract/primitive_infer_map.h" + +namespace mindspore { +namespace ops { +namespace { +abstract::ShapePtr InferShape(const PrimitivePtr &primitive, const std::vector &input_args) { + MS_EXCEPTION_IF_NULL(primitive); + auto prim_name = primitive->name(); + CheckAndConvertUtils::CheckInteger("input numbers", input_args.size(), kEqual, 1, prim_name); + auto x_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[0]->BuildShape())[kShape]; + auto axis = GetValue(primitive->GetAttr(kAxis)); + auto x_rank = SizeToLong(x_shape.size()); + CheckAndConvertUtils::CheckInRange("axis value", axis, kIncludeLeft, {-x_rank, x_rank}, prim_name); + return std::make_shared(x_shape); +} + +TypePtr InferType(const PrimitivePtr &prim, const std::vector &input_args) { + for (const auto &item : input_args) { + MS_EXCEPTION_IF_NULL(item); + } + const std::set valid_types = {kFloat32, kFloat16, kInt32, kUInt32, kInt8, kUInt8}; + auto infer_type = input_args[0]->BuildType(); + return CheckAndConvertUtils::CheckTensorTypeValid("x type", infer_type, valid_types, prim->name()); +} +} // namespace + +AbstractBasePtr RollInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive, + const std::vector &input_args) { + MS_EXCEPTION_IF_NULL(primitive); + return abstract::MakeAbstract(InferShape(primitive, input_args), InferType(primitive, input_args)); +} +REGISTER_PRIMITIVE_EVAL_IMPL(Roll, prim::kPrimRoll, RollInfer, nullptr, true); +} // namespace ops +} // namespace mindspore diff --git a/mindspore/core/ops/roll.h b/mindspore/core/ops/roll.h new file mode 100644 index 00000000000..8179f67d440 --- /dev/null +++ b/mindspore/core/ops/roll.h @@ -0,0 +1,41 @@ +/** + * Copyright 2021 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. + */ + +#ifndef MINDSPORE_CORE_OPS_ROLL_H_ +#define MINDSPORE_CORE_OPS_ROLL_H_ +#include +#include + +#include "ops/primitive_c.h" +#include "abstract/abstract_value.h" +#include "utils/check_convert_utils.h" + +namespace mindspore { +namespace ops { +constexpr auto kNameRoll = "Roll"; +class Roll : public PrimitiveC { + public: + Roll() : PrimitiveC(kNameRoll) { InitIOName({"input_x"}, {"output"}); } + ~Roll() = default; + MS_DECLARE_PARENT(Roll, PrimitiveC); +}; + +AbstractBasePtr RollInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive, + const std::vector &input_args); +using PrimRollPtr = std::shared_ptr; +} // namespace ops +} // namespace mindspore +#endif // MINDSPORE_CORE_OPS_ROLL_H_ diff --git a/mindspore/nn/layer/basic.py b/mindspore/nn/layer/basic.py index 11ee7cfae41..1c85b3843a3 100644 --- a/mindspore/nn/layer/basic.py +++ b/mindspore/nn/layer/basic.py @@ -33,7 +33,7 @@ from ..cell import Cell from .activation import get_activation __all__ = ['Dropout', 'Flatten', 'Dense', 'ClipByNorm', 'Norm', 'OneHot', 'Pad', 'Unfold', - 'Tril', 'Triu', 'ResizeBilinear', 'MatrixDiag', 'MatrixDiagPart', 'MatrixSetDiag', 'L1Regularizer'] + 'Tril', 'Triu', 'ResizeBilinear', 'MatrixDiag', 'MatrixDiagPart', 'MatrixSetDiag', 'L1Regularizer', 'Roll'] class L1Regularizer(Cell): @@ -1355,3 +1355,88 @@ class MatrixSetDiag(Cell): assist = _get_matrix_diag_part_assist(x_shape, x_dtype) out_matrix_set_diag = self.matrix_set_diag(input_x, diagonal, assist) return out_matrix_set_diag + + +@constexpr +def _check_input_dim(axis, dim, cls_name): + Validator.check_int_range(axis, -dim, dim, Rel.INC_LEFT, 'axis', cls_name) + + +class Roll(Cell): + """ + Rolls the elements of a tensor along an axis. + + The elements are shifted positively (towards larger indices) by the offset of `shift` along the dimension of `axis`. + Negative `shift` values will shift elements in the opposite direction. Elements that roll passed the last position + will wrap around to the first and vice versa. Multiple shifts along multiple axes may be specified. + + Args: + shift (Union[list(int), tuple(int), int]): Specifies the number of places by which elements are shifted + positively (towards larger indices) along the specified dimension. Negative shifts will roll the elements + in the opposite direction. + axis (Union[list(int), tuple(int), int]): Specifies the dimension indexes of shape to be rolled. + + Inputs: + - **input_x** (Tensor) - Input tensor. + + Outputs: + Tensor, has the same shape and type as `input_x`. + + Raises: + TypeError: If `shift` is not an int, a tuple or a list. + TypeError: If `axis` is not an int, a tuple or a list. + TypeError: If element of `shift` is not an int. + TypeError: If element of `axis` is not an int. + ValueError: If axis is out of the range [-len(input_x.shape), len(input_x.shape)). + ValueError: If length of shape of `shift` is not equal to length of shape of `axis`. + + Supported Platforms: + ``Ascend`` + + Examples: + >>> input_x = Tensor(np.array([0, 1, 2, 3, 4]).astype(np.float32)) + >>> op = nn.Roll(shift=2, axis=0) + >>> output = op(input_x) + >>> print(output) + [3. 4. 0. 1. 2.] + >>> input_x = Tensor(np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]).astype(np.float32)) + >>> op = nn.Roll(shift=[1, -2], axis=[0, 1]) + >>> output = op(input_x) + >>> print(output) + [[7. 8. 9. 5. 6.] + [2. 3. 4. 0. 1.]] + """ + + def __init__(self, shift, axis): + """Initialize Roll""" + super(Roll, self).__init__() + Validator.check_value_type("shift", shift, [int, tuple, list], self.cls_name) + Validator.check_value_type("axis", axis, [int, tuple, list], self.cls_name) + self.shape_op = P.Shape() + self.shift = shift + self.axis = axis + self.op_list = [] + + if not isinstance(self.axis, (list, tuple)): + self.op_list.append((inner.Roll(shift=self.shift, axis=0), self.axis)) + else: + if len(self.shift) != len(self.axis): + raise ValueError('The shape of shift and the shape of axis must be the same.') + for idx, _ in enumerate(self.axis): + self.op_list.append((inner.Roll(shift=self.shift[idx], axis=0), self.axis[idx])) + + def construct(self, input_x): + dim = len(self.shape_op(input_x)) + for single_op_roll, single_axis in self.op_list: + _check_input_dim(single_axis, dim, self.cls_name) + if single_axis < 0: + single_axis += dim + transpose_perm = [] + for i in range(dim): + transpose_perm.append(i) + transpose_perm[0], transpose_perm[single_axis] = single_axis, 0 + + input_x = input_x.transpose(transpose_perm) + input_x = single_op_roll(input_x) + input_x = input_x.transpose(transpose_perm) + return input_x diff --git a/mindspore/ops/_grad_experimental/grad_inner_ops.py b/mindspore/ops/_grad_experimental/grad_inner_ops.py index ff84e8ffd65..be38eefaa61 100644 --- a/mindspore/ops/_grad_experimental/grad_inner_ops.py +++ b/mindspore/ops/_grad_experimental/grad_inner_ops.py @@ -31,3 +31,17 @@ def get_bprop_tensor_copy_slices(self): return x_grad, update_grad, zeros_like(begin), zeros_like(end), zeros_like(stride) return bprop + + +@bprop_getters.register(inner.Roll) +def get_bprop_roll(self): + """Generate bprop for Roll""" + shift = self.shift + axis = self.axis + roll_grad = inner.Roll(-shift, axis) + + def bprop(x_input, out, dout): + dx = roll_grad(dout) + return (dx,) + + return bprop diff --git a/mindspore/ops/_op_impl/tbe/__init__.py b/mindspore/ops/_op_impl/tbe/__init__.py index bbe8f0cee3e..4024bf2b5e6 100644 --- a/mindspore/ops/_op_impl/tbe/__init__.py +++ b/mindspore/ops/_op_impl/tbe/__init__.py @@ -391,6 +391,7 @@ from .not_equal_ds import _not_ds_equal_tbe from .reciprocal_ds import _reciprocal_ds_tbe from .ctc_loss_v2 import _ctc_loss_v2_tbe from .ctc_loss_v2_grad import _ctc_loss_v2_grad_tbe +from .roll import _roll_tbe from .soft_shrink import _soft_shrink_tbe from .soft_shrink_grad import _soft_shrink_grad_tbe from .hsigmoid_grad import _hsigmoid_grad_tbe diff --git a/mindspore/ops/_op_impl/tbe/roll.py b/mindspore/ops/_op_impl/tbe/roll.py new file mode 100644 index 00000000000..7fccbe8a808 --- /dev/null +++ b/mindspore/ops/_op_impl/tbe/roll.py @@ -0,0 +1,42 @@ +# Copyright 2021 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. +# ============================================================================ + +"""Roll op""" +from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType + +roll_op_info = TBERegOp("Roll") \ + .fusion_type("OPAQUE") \ + .async_flag(False) \ + .binfile_name("roll.so") \ + .compute_cost(10) \ + .kernel_name("roll") \ + .partial_flag(True) \ + .attr("shift", "required", "listInt", "all") \ + .attr("axis", "optional", "listInt", "all") \ + .input(0, "input_x", False, "required", "all") \ + .output(0, "output", False, "required", "all") \ + .dtype_format(DataType.F32_Default, DataType.F32_Default) \ + .dtype_format(DataType.F16_Default, DataType.F16_Default) \ + .dtype_format(DataType.I32_Default, DataType.I32_Default) \ + .dtype_format(DataType.U32_Default, DataType.U32_Default) \ + .dtype_format(DataType.I8_Default, DataType.I8_Default) \ + .dtype_format(DataType.U8_Default, DataType.U8_Default) \ + .get_op_info() + + +@op_info_register(roll_op_info) +def _roll_tbe(): + """Roll TBE register""" + return diff --git a/mindspore/ops/operations/_inner_ops.py b/mindspore/ops/operations/_inner_ops.py index 07acdef27f4..418b3385ca6 100755 --- a/mindspore/ops/operations/_inner_ops.py +++ b/mindspore/ops/operations/_inner_ops.py @@ -1220,3 +1220,69 @@ class TensorCopySlices(Primitive): def __init__(self): """Initialize TensorScatterUpdate""" self.init_prim_io_names(inputs=['x', 'value', 'begin', 'end', 'strides'], outputs=['y']) + + +class Roll(Primitive): + """ + Rolls the elements of a tensor along an axis. + + The elements are shifted positively (towards larger indices) by the offset of `shift` along the dimension of `axis`. + Negative `shift` values will shift elements in the opposite direction. Elements that roll passed the last position + will wrap around to the first and vice versa. Multiple shifts along multiple axes may be specified. + + Note: + This inner operation is valid only if the axis is equal to 0. If the shift and the axis are tuples or lists, + this inner operation is valid only for the first pair of elements. + + Args: + shift (Union[list(int), tuple(int), int]): Specifies the number of places by which elements are shifted + positively (towards larger indices) along the specified dimension. Negative shifts will roll the elements + in the opposite direction. + axis (Union[list(int), tuple(int), int]): Specifies the dimension indexes of shape to be rolled. The value is + forced to be zero in this operation. + + Inputs: + - **input_x** (Tensor) - Input tensor. + + Outputs: + Tensor, has the same shape and type as `input_x`. + + Raises: + TypeError: If `shift` is not an int, a tuple or a list. + TypeError: If `axis` is not an int, a tuple or a list. + TypeError: If element of `shift` is not an int. + TypeError: If element of `axis` is not an int. + ValueError: If axis is not equal to 0. + ValueError: If shape of `shift` is not equal to 1. + ValueError: If shape of `axis` is not equal to 1. + + Supported Platforms: + ``Ascend`` + + Examples: + >>> from mindspore.ops.operations import _inner_ops as inner + >>> input_x = Tensor(np.array([0, 1, 2, 3, 4]).astype(np.float32)) + >>> op = inner.Roll(shift=2, axis=0) + >>> output = op(input_x) + >>> print(output) + [3. 4. 0. 1. 2.] + >>> input_x = Tensor(np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]).astype(np.float32)) + >>> op = inner.Roll(shift=-1, axis=0) + >>> output = op(input_x) + >>> print(output) + [[5. 6. 7. 8. 9.] + [0. 1. 2. 3. 4.]] + """ + + @prim_attr_register + def __init__(self, shift, axis): + """Initialize Roll""" + validator.check_value_type("shift", shift, [int, tuple, list], self.name) + validator.check_value_type("axis", axis, [int, tuple, list], self.name) + if isinstance(shift, (tuple, list)) and isinstance(axis, (tuple, list)): + validator.check_equal_int(len(shift), 1, "shift size", self.name) + validator.check_equal_int(len(axis), 1, "shift size", self.name) + validator.check_equal_int(axis[0], 0, "axis", self.name) + elif isinstance(shift, int) and isinstance(axis, int): + validator.check_equal_int(axis, 0, "axis", self.name) + self.init_prim_io_names(inputs=['input_x'], outputs=['output']) diff --git a/tests/ut/python/ops/test_ops.py b/tests/ut/python/ops/test_ops.py index bbff59c946f..9adb72044ff 100755 --- a/tests/ut/python/ops/test_ops.py +++ b/tests/ut/python/ops/test_ops.py @@ -2184,6 +2184,10 @@ test_case_nn_ops = [ Tensor(np.zeros((1, 1, 2, 2)), mstype.uint16)], 'desc_bprop': [], 'skip': ['backward']}), + ('Roll', { + 'block': nn.Roll(shift=[1, -2], axis=[0, 1]), + 'desc_inputs': [Tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]], mstype.float32)], + 'desc_bprop': [Tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]], mstype.float32)]}), ('SoftShrink', { 'block': P.SoftShrink(), 'desc_inputs': [Tensor(np.array([[0.5297, 0.7871, 1.1754], [0.7836, 0.6218, -1.1542]]), mstype.float32)],