From 6098c1da5e3167ff538a4996e0b810f73d7581f6 Mon Sep 17 00:00:00 2001 From: zorax <707120797@qq.com> Date: Wed, 21 Jul 2021 21:14:12 +0800 Subject: [PATCH] [feat][assistant][I40FFO] add new Ascend operator ApplyKerasMomentum --- mindspore/core/base/core_ops.h | 1 + mindspore/core/ops/apply_keras_momentum.cc | 90 ++++++++++++++++++ mindspore/core/ops/apply_keras_momentum.h | 47 +++++++++ mindspore/ops/_op_impl/tbe/__init__.py | 1 + .../ops/_op_impl/tbe/apply_keras_momentum.py | 49 ++++++++++ mindspore/ops/operations/__init__.py | 3 +- mindspore/ops/operations/nn_ops.py | 95 +++++++++++++++++++ tests/ut/python/ops/test_ops.py | 18 ++++ 8 files changed, 303 insertions(+), 1 deletion(-) create mode 100644 mindspore/core/ops/apply_keras_momentum.cc create mode 100644 mindspore/core/ops/apply_keras_momentum.h create mode 100644 mindspore/ops/_op_impl/tbe/apply_keras_momentum.py diff --git a/mindspore/core/base/core_ops.h b/mindspore/core/base/core_ops.h index 59525e7ed24..167f1f61566 100644 --- a/mindspore/core/base/core_ops.h +++ b/mindspore/core/base/core_ops.h @@ -412,6 +412,7 @@ inline const PrimitivePtr kPrimHShrink = std::make_shared("HShrink"); inline const PrimitivePtr kPrimHShrinkGrad = std::make_shared("HShrinkGrad"); inline const PrimitivePtr kPrimApplyAdagradDA = std::make_shared("ApplyAdagradDA"); inline const PrimitivePtr kPrimSparseApplyRMSProp = std::make_shared("SparseApplyRMSProp"); +inline const PrimitivePtr kPrimApplyKerasMomentum = std::make_shared("ApplyKerasMomentum"); // Comm ops inline const PrimitivePtr kPrimMirror = std::make_shared("_MirrorOperator"); diff --git a/mindspore/core/ops/apply_keras_momentum.cc b/mindspore/core/ops/apply_keras_momentum.cc new file mode 100644 index 00000000000..38a68176d6b --- /dev/null +++ b/mindspore/core/ops/apply_keras_momentum.cc @@ -0,0 +1,90 @@ +/** + * 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/apply_keras_momentum.h" + +#include +#include + +#include "ops/op_utils.h" +#include "abstract/primitive_infer_map.h" +#include "utils/tensor_construct_utils.h" + +namespace mindspore { +namespace ops { +namespace { +abstract::TupleShapePtr InferShape(const PrimitivePtr &primitive, const std::vector &input_args) { + MS_EXCEPTION_IF_NULL(primitive); + auto prim_name = primitive->name(); + (void)CheckAndConvertUtils::CheckInteger("input number", SizeToLong(input_args.size()), kEqual, 5, prim_name); + for (const auto &item : input_args) { + MS_EXCEPTION_IF_NULL(item); + } + auto var_shape = input_args[0]->BuildShape(); + auto accum_shape = input_args[1]->BuildShape(); + auto lr_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[2]->GetShapeTrack())[kShape]; + auto grad_shape = input_args[3]->BuildShape(); + auto momentum_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[4]->GetShapeTrack())[kShape]; + // lr, momentum must be scalar + (void)CheckAndConvertUtils::CheckInteger("lr_shape size", lr_shape.size(), kEqual, 0, prim_name); + (void)CheckAndConvertUtils::CheckInteger("momentum_shape size", momentum_shape.size(), kEqual, 0, prim_name); + // var, accum and grad must have the same shape + std::map same_shape_args_map; + same_shape_args_map.insert({"accum", accum_shape}); + same_shape_args_map.insert({"grad", grad_shape}); + for (auto &elem : same_shape_args_map) { + if (*elem.second != *var_shape) { + MS_EXCEPTION(ValueError) << prim_name << " evaluator arg " << elem.first << " shape " << elem.second->ToString() + << " are not consistent with var shape " << var_shape->ToString(); + } + } + return std::make_shared(std::vector{var_shape, accum_shape}); +} + +TuplePtr InferType(const PrimitivePtr &prim, const std::vector &input_args) { + MS_EXCEPTION_IF_NULL(prim); + auto prim_name = prim->name(); + (void)CheckAndConvertUtils::CheckInteger("input number", SizeToLong(input_args.size()), kEqual, 5, prim_name); + for (const auto &item : input_args) { + MS_EXCEPTION_IF_NULL(item); + } + auto var_type = input_args[0]->BuildType(); + auto accum_type = input_args[1]->BuildType(); + auto lr_type = input_args[2]->BuildType(); + auto grad_type = input_args[3]->BuildType(); + auto momentum_type = input_args[4]->BuildType(); + const std::set valid_types = {kFloat16, kFloat32}; + // var, accum and grad must have the same type + std::map args; + args.insert({"var", var_type}); + args.insert({"accum", accum_type}); + args.insert({"grad", grad_type}); + (void)CheckAndConvertUtils::CheckTensorTypeSame(args, valid_types, prim_name); + // lr, momentum type must be valid + (void)CheckAndConvertUtils::CheckTensorTypeValid("lr_dtype", lr_type, valid_types, prim_name); + (void)CheckAndConvertUtils::CheckTensorTypeValid("momentum_dtype", momentum_type, valid_types, prim_name); + return std::make_shared(std::vector{var_type, accum_type}); +} +} // namespace + +AbstractBasePtr ApplyKerasMomentumInfer(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(ApplyKerasMomentum, prim::kPrimApplyKerasMomentum, ApplyKerasMomentumInfer, nullptr, true); +} // namespace ops +} // namespace mindspore diff --git a/mindspore/core/ops/apply_keras_momentum.h b/mindspore/core/ops/apply_keras_momentum.h new file mode 100644 index 00000000000..75d4c74df49 --- /dev/null +++ b/mindspore/core/ops/apply_keras_momentum.h @@ -0,0 +1,47 @@ +/** + * 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_APPLY_KERAS_MOMENTUM_H_ +#define MINDSPORE_CORE_OPS_APPLY_KERAS_MOMENTUM_H_ + +#include +#include +#include +#include + +#include "ops/primitive_c.h" +#include "abstract/abstract_value.h" +#include "utils/check_convert_utils.h" + +namespace mindspore { +namespace ops { +constexpr auto kNameApplyKerasMomentum = "ApplyKerasMomentum"; +class MS_CORE_API ApplyKerasMomentum : public PrimitiveC { + public: + ApplyKerasMomentum() : PrimitiveC(kNameApplyKerasMomentum) { + InitIOName({"var", "accum", "lr", "grad", "momentum"}, {"var", "accum"}); + } + ~ApplyKerasMomentum() = default; + MS_DECLARE_PARENT(ApplyKerasMomentum, PrimitiveC); +}; + +AbstractBasePtr ApplyKerasMomentumInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive, + const std::vector &input_args); + +using PrimApplyKerasMomentumPtr = std::shared_ptr; +} // namespace ops +} // namespace mindspore +#endif // MINDSPORE_CORE_OPS_APPLY_KERAS_MOMENTUM_H_ diff --git a/mindspore/ops/_op_impl/tbe/__init__.py b/mindspore/ops/_op_impl/tbe/__init__.py index c8dcc57e42a..6772652cef0 100644 --- a/mindspore/ops/_op_impl/tbe/__init__.py +++ b/mindspore/ops/_op_impl/tbe/__init__.py @@ -28,6 +28,7 @@ from .add_n import _add_n_tbe from .add_n_ds import _add_n_ds_tbe from .accumulate_n_v2 import _accumulate_n_v2_tbe from .apply_ftrl import _apply_ftrl_tbe +from .apply_keras_momentum import _apply_keras_momentum_tbe from .apply_momentum import _apply_momentum_tbe from .apply_adam import _apply_adam_tbe from .apply_ada_max import _apply_ada_max_tbe diff --git a/mindspore/ops/_op_impl/tbe/apply_keras_momentum.py b/mindspore/ops/_op_impl/tbe/apply_keras_momentum.py new file mode 100644 index 00000000000..363dbbc2e40 --- /dev/null +++ b/mindspore/ops/_op_impl/tbe/apply_keras_momentum.py @@ -0,0 +1,49 @@ + # 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. +# ============================================================================ + +"""ApplyKerasMomentum op""" +from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType + +apply_keras_momentum_op_info = TBERegOp("ApplyKerasMomentum") \ + .fusion_type("OPAQUE") \ + .async_flag(False) \ + .binfile_name("apply_keras_momentum_d.so") \ + .compute_cost(10) \ + .kernel_name("apply_keras_momentum_d") \ + .partial_flag(True) \ + .attr("use_locking", "optional", "bool", "true,false", "false") \ + .attr("use_nesterov", "optional", "bool", "true,false", "false") \ + .input(0, "var", False, "required", "all") \ + .input(1, "accum", False, "required", "all") \ + .input(2, "lr", False, "required", "all") \ + .input(3, "grad", False, "required", "all") \ + .input(4, "momentum", False, "required", "all") \ + .output(0, "var", False, "required", "all") \ + .output(1, "accum", False, "required", "all") \ + .dtype_format(DataType.F32_5HD, DataType.F32_5HD, DataType.F32_Default, DataType.F32_5HD, + DataType.F32_Default, DataType.F32_5HD, DataType.F32_5HD) \ + .dtype_format(DataType.F32_C1HWNCoC0, DataType.F32_C1HWNCoC0, DataType.F32_Default, DataType.F32_C1HWNCoC0, + DataType.F32_Default, DataType.F32_C1HWNCoC0, DataType.F32_C1HWNCoC0) \ + .dtype_format(DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, + DataType.F32_Default, DataType.F32_Default, DataType.F32_Default) \ + .dtype_format(DataType.F32_FracZ, DataType.F32_FracZ, DataType.F32_Default, DataType.F32_Default, + DataType.F32_Default, DataType.F32_FracZ, DataType.F32_FracZ) \ + .get_op_info() + +@op_info_register(apply_keras_momentum_op_info) +def _apply_keras_momentum_tbe(): + """ApplyKerasMomentum TBE register""" + return + \ No newline at end of file diff --git a/mindspore/ops/operations/__init__.py b/mindspore/ops/operations/__init__.py index 2003695e96e..74f7af0b86e 100644 --- a/mindspore/ops/operations/__init__.py +++ b/mindspore/ops/operations/__init__.py @@ -76,7 +76,7 @@ from .nn_ops import (LSTM, SGD, Adam, AdamWeightDecay, FusedSparseAdam, FusedSpa MaxPool, DataFormatDimMap, AvgPool, Conv2DBackpropInput, ComputeAccidentalHits, MaxPoolWithArgmax, OneHot, Pad, MirrorPad, Mish, PReLU, ReLU, ReLU6, ReLUV2, HSwish, HSigmoid, - ResizeBilinear, Sigmoid, SeLU, HShrink, + ResizeBilinear, Sigmoid, SeLU, HShrink, ApplyKerasMomentum, SigmoidCrossEntropyWithLogits, NLLLoss, BCEWithLogitsLoss, SmoothL1Loss, SoftMarginLoss, Softmax, Softsign, Softplus, LRN, RNNTLoss, DynamicRNN, DynamicGRUV2, SoftmaxCrossEntropyWithLogits, ROIAlign, @@ -185,6 +185,7 @@ __all__ = [ 'NLLLoss', 'SGD', 'ApplyMomentum', + "ApplyKerasMomentum", 'FusedWeightScaleApplyMomentum', 'ExpandDims', 'Cast', diff --git a/mindspore/ops/operations/nn_ops.py b/mindspore/ops/operations/nn_ops.py index 364e9e1e2e2..cccc206a8d8 100755 --- a/mindspore/ops/operations/nn_ops.py +++ b/mindspore/ops/operations/nn_ops.py @@ -9301,3 +9301,98 @@ class SparseApplyRMSProp(Primitive): self.epsilon = validator.check_number("epsilon", epsilon, 0.0, Rel.GT, self.name) self.momentum = validator.check_number("momentum", momentum, 0.0, Rel.GE, self.name) self.rho = validator.check_float_range(rho, 0.0, 1.0, Rel.INC_BOTH, "rho", self.name) + + +class ApplyKerasMomentum(Primitive): + r""" + Update `var` according to the momentum scheme. + + .. math:: + \begin{array}{ll} \\ + accum = accum * momentum - grad * lr \\ + var = + \begin{cases} + var + accum * momentum - grad * lr, &\text{if use_nesterov} \\ + var + accum, &\text{else} + \end{cases} + \end{array} + + Refer to the paper `On the importance of initialization and momentum in deep + learning `_ for more details. + + Inputs of `var`, `accum` and `grad` comply with the implicit type conversion rules + to make the data types consistent. + If they have different data types, lower priority data type will be converted to + relatively highest priority data type. + RuntimeError exception will be thrown when the data type conversion of Parameter is required. + + Args: + use_locking (bool): If `True`, updating of the `var` and `accum` tensors will be protected by a lock; + Otherwise the behavior is undefined, but may exhibit less contention. Default: False. + use_nesterov (bool): If `True`, the tensor passed to compute grad will be var + momentum * accum, + so in the end, the var you get is actually var + momentum * accum. Default: False. + + Inputs: + - **var** (Parameter) - Variable to be updated. With float16 or float32 data type. + - **accum** (Parameter) - Must have the same shape and type as `var`. With float16 or float32 data type. + - **lr** (Union[Number, Tensor]) - Scaling factor. Must be a scalar. With float16 or float32 data type. + - **grad** (Tensor) - The gradient. Must have the same shape and type as `var`. + With float16 or float32 data type. + - **momentum** (Union[Number, Tensor]) - Momentum. Must be a scalar. With float16 or float32 data type. + + Outputs: + Tuple of 2 Tensors, the updated parameters. + + - **var** (Tensor) - The same shape and data type as `var`. + - **accum** (Tensor) - The same shape and data type as `accum`. + + Raises: + TypeError: If the use_locking or use_nesterov is not a bool. + TypeError: If `var` or `accum` is not a Parameter. + TypeError: If `lr` is neither a Number nor a Tensor. + TypeError: If `grad` is not a Tensor. + TypeError: If `momentum` is neither a Number nor a Tensor. + TypeError: If dtype of `var`, `accum`, `lr`, `grad`, `momentum` is neither float16 nor float32. + ValueError: If `accum` or `grad` doesn't have the same shape as `var`. + ValueError: If the shape size of `lr`, `momentum` is not 0. + + Supported Platforms: + ``Ascend`` + + Examples: + >>> class ApplyKerasMomentumNet(nn.Cell): + ... def __init__(self, use_locking=False, use_nesterov=False): + ... super(ApplyKerasMomentumNet, self).__init__() + ... self.apply_keras_momentum = P.ApplyKerasMomentum(use_locking, use_nesterov) + ... self.var = Parameter(Tensor(np.array([[0.2, 0.3], [0.1, 0.4]]).astype(np.float32)), name="var") + ... self.accum = Parameter(Tensor(np.array([[0.2, 0.3], [0.1, 0.4]]).astype(np.float32)), name="accum") + ... def construct(self, lr, grad, momentum): + ... out = self.apply_keras_momentum(self.var, self.accum, lr, grad, momentum) + ... return out + ... + >>> net = ApplyKerasMomentumNet() + >>> lr = Tensor(0.001, mstype.float32) + >>> grad = Tensor(np.array([[0.3, 0.2], [0.4, 0.1]]).astype(np.float32)) + >>> momentum = Tensor(0.99, mstype.float32) + >>> output = net(lr, grad, momentum) + >>> print(output) + (Tensor(shape=[2, 2], dtype=Float32, value= + [[ 3.97700012e-01, 5.96800029e-01], + [ 1.98599994e-01, 7.95899987e-01]]), Tensor(shape=[2, 2], dtype=Float32, value= + [[ 1.97699994e-01, 2.96800017e-01], + [ 9.86000001e-02, 3.95900011e-01]])) + """ + + __mindspore_signature__ = ( + sig.make_sig('var', sig.sig_rw.RW_WRITE, dtype=sig.sig_dtype.T), + sig.make_sig('accum', sig.sig_rw.RW_WRITE, dtype=sig.sig_dtype.T), + sig.make_sig('lr', dtype=sig.sig_dtype.T1), + sig.make_sig('grad', dtype=sig.sig_dtype.T), + sig.make_sig('momentum', dtype=sig.sig_dtype.T2) + ) + + @prim_attr_register + def __init__(self, use_locking=False, use_nesterov=False): + """Initialize ApplyKerasMomentum""" + validator.check_value_type("use_locking", use_locking, [bool], self.name) + validator.check_value_type("use_nesterov", use_nesterov, [bool], self.name) diff --git a/tests/ut/python/ops/test_ops.py b/tests/ut/python/ops/test_ops.py index b3e1c81070e..882440ad618 100755 --- a/tests/ut/python/ops/test_ops.py +++ b/tests/ut/python/ops/test_ops.py @@ -1042,6 +1042,18 @@ class SparseApplyRMSPropNet(nn.Cell): out = self.sparse_apply_r_m_s_prop(self.var, self.ms, self.mom, lr, grad, indices) return out + +class ApplyKerasMomentumNet(nn.Cell): + def __init__(self, use_locking=False, use_nesterov=False): + super(ApplyKerasMomentumNet, self).__init__() + self.apply_keras_momentum = P.ApplyKerasMomentum(use_locking, use_nesterov) + self.var = Parameter(Tensor(np.array([[0.2, 0.3], [0.1, 0.4]]).astype(np.float32)), name="var") + self.accum = Parameter(Tensor(np.array([[0.2, 0.3], [0.1, 0.4]]).astype(np.float32)), name="accum") + + def construct(self, lr, grad, momentum): + out = self.apply_keras_momentum(self.var, self.accum, lr, grad, momentum) + return out + test_case_math_ops = [ ('BitwiseAnd', { 'block': P.BitwiseAnd(), @@ -2306,6 +2318,12 @@ test_case_nn_ops = [ Tensor(np.array([[0.3, 0.7], [0.1, 0.8]]).astype(np.float32)), Tensor(np.array([0, 1], dtype=np.int32))], 'skip': ['backward']}), + ('ApplyKerasMomentum', { + 'block': ApplyKerasMomentumNet(), + 'desc_inputs': [Tensor(0.001, mstype.float32), + Tensor(np.array([[0.3, 0.2], [0.4, 0.1]]).astype(np.float32)), + Tensor(0.99, mstype.float32)], + 'skip': ['backward']}), ] test_case_array_ops = [