forked from mindspore-Ecosystem/mindspore
Add SparseApplyFtrl cpu kernel
This commit is contained in:
parent
06ee0296b8
commit
5d4b75838f
|
@ -525,5 +525,27 @@ std::string GetProcessor(const AnfNodePtr &anf_node) {
|
|||
}
|
||||
return device;
|
||||
}
|
||||
|
||||
bool IsSameShape(const std::vector<size_t> &shape_a, const std::vector<size_t> &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
|
||||
|
|
|
@ -82,6 +82,8 @@ bool ParseMetadata(const CNodePtr &kernel_node, const std::shared_ptr<const OpIn
|
|||
bool IsAtomicNode(const CNodePtr &kernel_node);
|
||||
void SaveJsonInfo(const std::string &json_name, const std::string &info);
|
||||
std::string GetProcessor(const AnfNodePtr &anf_node);
|
||||
bool IsSameShape(const std::vector<size_t> &shape_a, const std::vector<size_t> &shape_b);
|
||||
int Sign(float x);
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
||||
|
||||
|
|
|
@ -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<size_t> var_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 0);
|
||||
std::vector<size_t> accum_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 1);
|
||||
std::vector<size_t> linear_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 2);
|
||||
std::vector<size_t> grad_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 3);
|
||||
std::vector<size_t> 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<float>(kernel_node, "lr");
|
||||
if (lr_ <= 0) {
|
||||
MS_LOG(EXCEPTION) << "lr should be a positive scalar";
|
||||
}
|
||||
l1_ = AnfAlgo::GetNodeAttr<float>(kernel_node, "l1");
|
||||
if (l1_ < 0) {
|
||||
MS_LOG(EXCEPTION) << "l1 should be a non-negative scalar";
|
||||
}
|
||||
l2_ = AnfAlgo::GetNodeAttr<float>(kernel_node, "l2");
|
||||
if (l2_ < 0) {
|
||||
MS_LOG(EXCEPTION) << "l2 should be a non-negative scalar";
|
||||
}
|
||||
lr_power_ = AnfAlgo::GetNodeAttr<float>(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<kernel::AddressPtr> &inputs,
|
||||
const std::vector<kernel::AddressPtr> & /*workspace*/,
|
||||
const std::vector<kernel::AddressPtr> & /*outputs*/) {
|
||||
if (inputs.size() < kSparseApplyFtrlInputSize) {
|
||||
MS_LOG(EXCEPTION) << "error input output size!";
|
||||
}
|
||||
|
||||
auto var = reinterpret_cast<float *>(inputs[0]->addr);
|
||||
auto accum = reinterpret_cast<float *>(inputs[1]->addr);
|
||||
auto linear = reinterpret_cast<float *>(inputs[2]->addr);
|
||||
auto grad = reinterpret_cast<float *>(inputs[3]->addr);
|
||||
auto indices = reinterpret_cast<int *>(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
|
|
@ -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 <vector>
|
||||
#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<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
|
||||
const std::vector<AddressPtr> &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_
|
|
@ -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())
|
Loading…
Reference in New Issue