1 /**
2 * Copyright 2020 Huawei Technologies Co., Ltd
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 #include "src/litert/kernel/cpu/fp32_grad/adam.h"
18 #include <cmath>
19 #include <string>
20 #include "schema/model_generated.h"
21 #include "src/litert/kernel_registry.h"
22 #include "include/errorcode.h"
23 #include "nnacl/fp32/adam_fp32.h"
24 #include "plugin/device/cpu/kernel/nnacl/op_base.h"
25
26 using mindspore::kernel::KERNEL_ARCH;
27 using mindspore::lite::KernelRegistrar;
28 using mindspore::lite::RET_ERROR;
29 using mindspore::lite::RET_OK;
30 using mindspore::schema::PrimitiveType_Adam;
31
32 namespace mindspore::kernel {
33 constexpr static int kBeta1PowerIdx = 3;
34 constexpr static int kBeta2PowerIdx = 4;
35 constexpr static int kBeta1Idx = 6;
36 constexpr static int kBeta2Idx = 7;
37 constexpr static int kEpsilonIdx = 8;
38 constexpr static int kGradientIdx = 9;
39
ReSize()40 int AdamCPUKernel::ReSize() { return RET_OK; }
41
DoExecute(int task_id)42 int AdamCPUKernel::DoExecute(int task_id) {
43 CHECK_LESS_RETURN(in_tensors_.size(), DIMENSION_10D);
44 auto weight = reinterpret_cast<float *>(in_tensors_.at(kWeightIdx)->MutableData());
45 auto m = reinterpret_cast<float *>(in_tensors_.at(kMomentVector1stIdx)->MutableData());
46 auto v = reinterpret_cast<float *>(in_tensors_.at(kMomentVector2stIdx)->MutableData());
47 auto beta1_power = reinterpret_cast<float *>(in_tensors_.at(kBeta1PowerIdx)->MutableData());
48 auto beta2_power = reinterpret_cast<float *>(in_tensors_.at(kBeta2PowerIdx)->MutableData());
49 auto learning_rate = lr_;
50 auto beta1 = reinterpret_cast<float *>(in_tensors_.at(kBeta1Idx)->MutableData())[0];
51 auto beta2 = reinterpret_cast<float *>(in_tensors_.at(kBeta2Idx)->MutableData())[0];
52 auto eps = reinterpret_cast<float *>(in_tensors_.at(kEpsilonIdx)->MutableData())[0];
53 auto gradient = reinterpret_cast<float *>(in_tensors_.at(kGradientIdx)->MutableData());
54 int length = in_tensors_.at(kWeightIdx)->ElementsNum();
55 CHECK_NULL_RETURN(weight);
56 CHECK_NULL_RETURN(m);
57 CHECK_NULL_RETURN(v);
58 CHECK_NULL_RETURN(gradient);
59
60 int stride = UP_DIV(length, thread_count_);
61 int count = MSMIN(stride, length - stride * task_id);
62 int start = stride * task_id;
63 int end = start + count;
64
65 return DoAdam(m, v, gradient, weight, beta1, beta2, beta1_power, beta2_power, eps, learning_rate,
66 adam_param_->use_nesterov_, start, end);
67 }
68
AdamRun(void * cdata,int task_id,float lhs_scale,float rhs_scale)69 int AdamRun(void *cdata, int task_id, float lhs_scale, float rhs_scale) {
70 auto adam_kernel = reinterpret_cast<AdamCPUKernel *>(cdata);
71 CHECK_NULL_RETURN(adam_kernel);
72 auto error_code = RET_OK;
73 if (adam_kernel->get_optimizer_mode() == WeightUpdateMode::VIRTUAL_BATCH) {
74 error_code = adam_kernel->ExecuteVirtualBatch(task_id);
75 } else if (adam_kernel->get_optimizer_mode() == WeightUpdateMode::ACCUMULATE_GRADS) {
76 error_code = adam_kernel->ExecuteVirtualBatch(task_id);
77 } else {
78 error_code = adam_kernel->DoExecute(task_id);
79 }
80
81 if (error_code != RET_OK) {
82 MS_LOG(ERROR) << "Adam run error task_id[" << task_id << "] error_code[" << error_code << "]";
83 return RET_ERROR;
84 }
85 return RET_OK;
86 }
87
Run()88 int AdamCPUKernel::Run() {
89 int error_code = ParallelLaunch(this->ms_context_, AdamRun, this, thread_count_);
90 if (error_code != RET_OK) {
91 MS_LOG(ERROR) << "Adam function error error_code[" << error_code << "]";
92 return RET_ERROR;
93 }
94 return RET_OK;
95 }
96
Prepare()97 int AdamCPUKernel::Prepare() {
98 CHECK_NULL_RETURN(adam_param_);
99 auto ret = OptimizerKernel::Prepare();
100 if (ret != RET_OK) {
101 MS_LOG(ERROR) << "Failed to initialize Adam Kernel";
102 return RET_ERROR;
103 }
104 return RET_OK;
105 }
106
GetOptimizerParamsIdxs() const107 std::vector<int> AdamCPUKernel::GetOptimizerParamsIdxs() const {
108 std::vector<int> indices = {6, 7, 3, 4, 8};
109 return indices;
110 }
111
GetTrainableParamsIdxs() const112 std::vector<int> AdamCPUKernel::GetTrainableParamsIdxs() const {
113 std::vector<int> indices = {0, 1, 2, 3, 4, 5};
114 return indices;
115 }
116
OptimizerStep()117 int AdamCPUKernel::OptimizerStep() {
118 CHECK_LESS_RETURN(in_tensors_.size(), DIMENSION_10D - 1);
119 auto weight = reinterpret_cast<float *>(in_tensors_.at(kWeightIdx)->MutableData());
120 auto m = reinterpret_cast<float *>(in_tensors_.at(kMomentVector1stIdx)->MutableData());
121 auto v = reinterpret_cast<float *>(in_tensors_.at(kMomentVector2stIdx)->MutableData());
122 auto beta1_power = reinterpret_cast<float *>(in_tensors_.at(kBeta1PowerIdx)->MutableData());
123 auto beta2_power = reinterpret_cast<float *>(in_tensors_.at(kBeta2PowerIdx)->MutableData());
124 auto learning_rate = lr_;
125 auto beta1 = reinterpret_cast<float *>(in_tensors_.at(kBeta1Idx)->MutableData())[0];
126 auto beta2 = reinterpret_cast<float *>(in_tensors_.at(kBeta2Idx)->MutableData())[0];
127 auto eps = reinterpret_cast<float *>(in_tensors_.at(kEpsilonIdx)->MutableData())[0];
128 size_t length = static_cast<size_t>(in_tensors_.at(kWeightIdx)->ElementsNum());
129 CHECK_NULL_RETURN(weight);
130 CHECK_NULL_RETURN(m);
131 CHECK_NULL_RETURN(v);
132 CHECK_NULL_RETURN(beta1_power);
133 CHECK_NULL_RETURN(beta2_power);
134
135 int ret = RET_OK;
136 if (grad_sum_ != nullptr && valid_grad_sum_) {
137 size_t start = 0;
138 size_t end = length;
139 ret = DoAdam(m, v, grad_sum_, weight, beta1, beta2, beta1_power, beta2_power, eps, learning_rate,
140 adam_param_->use_nesterov_, start, end);
141 std::fill(grad_sum_, grad_sum_ + length, 0);
142 (void)OptimizerKernel::OptimizerStep();
143 }
144 return ret;
145 }
146
CpuAdamFp32KernelCreator(const std::vector<lite::Tensor * > & inputs,const std::vector<lite::Tensor * > & outputs,OpParameter * opParameter,const lite::InnerContext * ctx,const kernel::KernelKey & desc)147 kernel::LiteKernel *CpuAdamFp32KernelCreator(const std::vector<lite::Tensor *> &inputs,
148 const std::vector<lite::Tensor *> &outputs, OpParameter *opParameter,
149 const lite::InnerContext *ctx, const kernel::KernelKey &desc) {
150 MS_CHECK_TRUE_MSG(opParameter != nullptr, nullptr, "Op parameter is nullptr.");
151 MS_ASSERT(desc.type == schema::PrimitiveType_Adam);
152 auto *kernel = new (std::nothrow) AdamCPUKernel(opParameter, inputs, outputs, ctx);
153 if (kernel == nullptr) {
154 MS_LOG(ERROR) << "new AdamCPUKernel fail!";
155 free(opParameter);
156 return nullptr;
157 }
158 return kernel;
159 }
160
161 REG_KERNEL(kCPU, kNumberTypeFloat32, PrimitiveType_Adam, CpuAdamFp32KernelCreator)
162 } // namespace mindspore::kernel
163