diff --git a/mindspore/ccsrc/kernel/common_utils.cc b/mindspore/ccsrc/kernel/common_utils.cc index 2769e0c42a..924592232f 100644 --- a/mindspore/ccsrc/kernel/common_utils.cc +++ b/mindspore/ccsrc/kernel/common_utils.cc @@ -525,5 +525,27 @@ std::string GetProcessor(const AnfNodePtr &anf_node) { } return device; } + +bool IsSameShape(const std::vector &shape_a, const std::vector &shape_b) { + if (shape_a.size() != shape_b.size()) { + return false; + } + for (size_t i = 0; i < shape_a.size(); ++i) { + if (shape_a[i] != shape_b[i]) { + return false; + } + } + return true; +} + +int Sign(float x) { + if (x > 0) { + return 1; + } + if (x < 0) { + return -1; + } + return 0; +} } // namespace kernel } // namespace mindspore diff --git a/mindspore/ccsrc/kernel/common_utils.h b/mindspore/ccsrc/kernel/common_utils.h index 47fe96c4c9..2d5161b870 100644 --- a/mindspore/ccsrc/kernel/common_utils.h +++ b/mindspore/ccsrc/kernel/common_utils.h @@ -82,6 +82,8 @@ bool ParseMetadata(const CNodePtr &kernel_node, const std::shared_ptr &shape_a, const std::vector &shape_b); +int Sign(float x); } // namespace kernel } // namespace mindspore diff --git a/mindspore/ccsrc/kernel/cpu/sparse_apply_ftrl_cpu_kernel.cc b/mindspore/ccsrc/kernel/cpu/sparse_apply_ftrl_cpu_kernel.cc new file mode 100644 index 0000000000..b74d3792ff --- /dev/null +++ b/mindspore/ccsrc/kernel/cpu/sparse_apply_ftrl_cpu_kernel.cc @@ -0,0 +1,115 @@ +/** + * Copyright 2020 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 "kernel/cpu/sparse_apply_ftrl_cpu_kernel.h" +#include "kernel/common_utils.h" +#include "device/cpu/cpu_device_address.h" + +namespace mindspore { +namespace kernel { +namespace { +constexpr size_t kSparseApplyFtrlInputSize = 5; +} // namespace + +void SparseApplyFtrlCPUKernel::InitKernel(const CNodePtr &kernel_node) { + MS_EXCEPTION_IF_NULL(kernel_node); + std::vector var_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 0); + std::vector accum_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 1); + std::vector linear_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 2); + std::vector grad_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 3); + std::vector indices_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 4); + if (!IsSameShape(var_shape, accum_shape)) { + MS_LOG(EXCEPTION) << "var and accum should have the same shape"; + } + if (!IsSameShape(var_shape, linear_shape)) { + MS_LOG(EXCEPTION) << "var and linear should have the same shape"; + } + if (var_shape.empty()) { + MS_LOG(EXCEPTION) << "var must be at least 1D"; + } + var_first_dim_size_ = var_shape[0]; + for (size_t i = 1; i < var_shape.size(); ++i) { + if (var_shape[i] != grad_shape[i]) { + MS_LOG(EXCEPTION) << "The shape of var and grad must equal in dimension " << i; + } + var_outer_dim_size_ *= var_shape[i]; + } + if (indices_shape.size() != 1) { + MS_LOG(EXCEPTION) << "indices must be a 1D vector"; + } + indices_size_ = indices_shape[0]; + if (grad_shape[0] != indices_size_) { + MS_LOG(ERROR) << "The first dimension of grad shape must be equal to indices"; + } + lr_ = AnfAlgo::GetNodeAttr(kernel_node, "lr"); + if (lr_ <= 0) { + MS_LOG(EXCEPTION) << "lr should be a positive scalar"; + } + l1_ = AnfAlgo::GetNodeAttr(kernel_node, "l1"); + if (l1_ < 0) { + MS_LOG(EXCEPTION) << "l1 should be a non-negative scalar"; + } + l2_ = AnfAlgo::GetNodeAttr(kernel_node, "l2"); + if (l2_ < 0) { + MS_LOG(EXCEPTION) << "l2 should be a non-negative scalar"; + } + lr_power_ = AnfAlgo::GetNodeAttr(kernel_node, "lr_power"); + if (lr_power_ > 0) { + MS_LOG(EXCEPTION) << "lr_power should be a non-positive scalar"; + } +} + +bool SparseApplyFtrlCPUKernel::Launch(const std::vector &inputs, + const std::vector & /*workspace*/, + const std::vector & /*outputs*/) { + if (inputs.size() < kSparseApplyFtrlInputSize) { + MS_LOG(EXCEPTION) << "error input output size!"; + } + + auto var = reinterpret_cast(inputs[0]->addr); + auto accum = reinterpret_cast(inputs[1]->addr); + auto linear = reinterpret_cast(inputs[2]->addr); + auto grad = reinterpret_cast(inputs[3]->addr); + auto indices = reinterpret_cast(inputs[4]->addr); + + for (size_t i = 0; i < indices_size_; ++i) { + int index = indices[i]; + if ((size_t)index >= var_first_dim_size_) { + MS_LOG(EXCEPTION) << "Index " << index << " in indices is out of range"; + } + for (size_t j = var_outer_dim_size_ * index, k = var_outer_dim_size_ * i; j < var_outer_dim_size_ * (index + 1); + ++j, ++k) { + auto accum_new = accum[j] + grad[k] * grad[k]; + if (lr_power_ == -0.5) { + linear[j] += grad[k] - (sqrt(accum_new) - sqrt(accum[j])) / lr_ * var[j]; + } else { + linear[j] += grad[k] - (pow(accum_new, -lr_power_) - pow(accum[j], -lr_power_)) / lr_ * var[j]; + } + auto x = Sign(linear[j]) * l1_ - linear[j]; + float y; + if (lr_power_ == -0.5) { + y = sqrt(accum_new) / lr_ + 2 * l2_; + } else { + y = pow(accum_new, -lr_power_) / lr_ + 2 * l2_; + } + auto pre_shrink = x / y; + var[j] = abs(linear[j]) > l1_ ? pre_shrink : 0; + accum[j] = accum_new; + } + } + return true; +} +} // namespace kernel +} // namespace mindspore diff --git a/mindspore/ccsrc/kernel/cpu/sparse_apply_ftrl_cpu_kernel.h b/mindspore/ccsrc/kernel/cpu/sparse_apply_ftrl_cpu_kernel.h new file mode 100644 index 0000000000..6dc075c45f --- /dev/null +++ b/mindspore/ccsrc/kernel/cpu/sparse_apply_ftrl_cpu_kernel.h @@ -0,0 +1,59 @@ +/** + * Copyright 2020 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_CCSRC_KERNEL_CPU_SPARSE_APPLY_FTRL_CPU_KERNEL_H_ +#define MINDSPORE_CCSRC_KERNEL_CPU_SPARSE_APPLY_FTRL_CPU_KERNEL_H_ + +#include +#include "kernel/cpu/cpu_kernel.h" +#include "kernel/cpu/cpu_kernel_factory.h" + +namespace mindspore { +namespace kernel { +class SparseApplyFtrlCPUKernel : public CPUKernel { + public: + SparseApplyFtrlCPUKernel() = default; + ~SparseApplyFtrlCPUKernel() override = default; + + void InitKernel(const CNodePtr &kernel_node) override; + + bool Launch(const std::vector &inputs, const std::vector &workspace, + const std::vector &outputs) override; + + private: + size_t indices_size_{0}; + size_t var_first_dim_size_{0}; + size_t var_outer_dim_size_{1}; + float lr_{0}; + float l1_{0}; + float l2_{0}; + float lr_power_{0}; +}; + +MS_REG_CPU_KERNEL(SparseApplyFtrl, + KernelAttr() + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeInt32) + .AddOutputAttr(kNumberTypeFloat32) + .AddOutputAttr(kNumberTypeFloat32) + .AddOutputAttr(kNumberTypeFloat32), + SparseApplyFtrlCPUKernel); +} // namespace kernel +} // namespace mindspore + +#endif // MINDSPORE_CCSRC_KERNEL_CPU_SPARSE_APPLY_FTRL_CPU_KERNEL_H_ diff --git a/tests/st/ops/cpu/test_sparse_apply_ftrl_op.py b/tests/st/ops/cpu/test_sparse_apply_ftrl_op.py new file mode 100644 index 0000000000..cc1c3e3507 --- /dev/null +++ b/tests/st/ops/cpu/test_sparse_apply_ftrl_op.py @@ -0,0 +1,50 @@ +# Copyright 2020 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. +# ============================================================================ + +import numpy as np +import mindspore.context as context +import mindspore.nn as nn +from mindspore import Tensor +from mindspore.common.parameter import Parameter +from mindspore.ops import operations as P +import mindspore.common.dtype as mstype + + +class Net(nn.Cell): + def __init__(self): + super(Net, self).__init__() + self.sparse_apply_ftrl = P.SparseApplyFtrl(lr=0.001, l1=0.0, l2=0.0, lr_power=-0.5) + self.var = Parameter(Tensor(np.ones([3, 3, 3]).astype(np.float32)), name="var") + self.accum = Parameter(Tensor(np.ones([3, 3, 3]).astype(np.float32)), name="accum") + self.linear = Parameter(Tensor(np.ones([3, 3, 3]).astype(np.float32)), name="linear") + + def construct(self, grad, indices): + out = self.sparse_apply_ftrl(self.var, self.accum, self.linear, grad, indices) + return out + + +def test_net(): + gradient = Tensor(np.random.rand(3, 3, 3).astype(np.float32)) + indices = Tensor([0, 1, 2], mstype.int32) + + context.set_context(mode=context.GRAPH_MODE, device_target="CPU") + sparse_apply_ftrl = Net() + output = sparse_apply_ftrl(gradient, indices) + print(output[0].asnumpy()) + + context.set_context(mode=context.GRAPH_MODE, device_target="Ascend") + sparse_apply_ftrl = Net() + output = sparse_apply_ftrl(gradient, indices) + print(output[0].asnumpy())