Add SparseApplyFtrl cpu kernel

This commit is contained in:
yujianfeng 2020-06-01 15:28:31 +08:00
parent 06ee0296b8
commit 5d4b75838f
5 changed files with 248 additions and 0 deletions

View File

@ -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

View File

@ -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

View File

@ -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

View File

@ -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_

View File

@ -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())