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/runtime/kernel/arm/fp32/batchnorm_fp32.h"
18 #include "src/kernel_registry.h"
19
20 using mindspore::lite::KernelRegistrar;
21 using mindspore::lite::RET_ERROR;
22 using mindspore::lite::RET_NULL_PTR;
23 using mindspore::lite::RET_OK;
24 using mindspore::schema::PrimitiveType_BatchNorm;
25 namespace {
26 constexpr int kNumInput2 = 2;
27 }
28 namespace mindspore::kernel {
Init()29 int BatchnormCPUKernel::Init() {
30 CHECK_LESS_RETURN(in_tensors_.size(), DIMENSION_3D);
31 CHECK_LESS_RETURN(out_tensors_.size(), 1);
32 CHECK_NULL_RETURN(in_tensors_[0]);
33 CHECK_NULL_RETURN(in_tensors_[1]);
34 CHECK_NULL_RETURN(in_tensors_[kNumInput2]);
35 CHECK_NULL_RETURN(out_tensors_[0]);
36 CHECK_NULL_RETURN(op_parameter_);
37 if (!InferShapeDone()) {
38 return RET_OK;
39 }
40 return ReSize();
41 }
42
ReSize()43 int BatchnormCPUKernel::ReSize() {
44 FreeMeanAndVariance();
45 auto status = FillParam();
46 if (status != RET_OK) {
47 return RET_ERROR;
48 }
49 return InitConstTensor();
50 }
51
FreeMeanAndVariance()52 void BatchnormCPUKernel::FreeMeanAndVariance() {
53 if (mean_ != nullptr) {
54 free(mean_);
55 mean_ = nullptr;
56 }
57 if (variance_ != nullptr) {
58 free(variance_);
59 variance_ = nullptr;
60 }
61 }
62
FillParam()63 int BatchnormCPUKernel::FillParam() {
64 auto input_shapes = in_tensors_.at(0)->shape();
65 auto n_dim = input_shapes.size();
66 auto param = reinterpret_cast<BatchNormParameter *>(op_parameter_);
67 CHECK_LESS_RETURN(n_dim - 1, 0);
68 param->channel_ = input_shapes[n_dim - 1];
69 param->unit_ = 1;
70 for (size_t i = 0; i < n_dim - 1; i++) {
71 param->unit_ *= input_shapes[i];
72 }
73 if (default_momentum_ < 0.0f) {
74 default_momentum_ = param->momentum_;
75 }
76 return RET_OK;
77 }
78
InitConstTensor()79 int BatchnormCPUKernel::InitConstTensor() {
80 CHECK_LESS_RETURN(MAX_MALLOC_SIZE, in_tensors_.at(1)->Size());
81 CHECK_LESS_RETURN(MAX_MALLOC_SIZE, in_tensors_.at(kNumInput2)->Size());
82 mean_ = malloc(in_tensors_.at(SECOND_INPUT)->Size());
83 variance_ = malloc(in_tensors_.at(THIRD_INPUT)->Size());
84 if (mean_ == nullptr || variance_ == nullptr) {
85 MS_LOG(ERROR) << "Memory allocation failed";
86 FreeMeanAndVariance();
87 return RET_ERROR;
88 }
89 auto in_tensor_mean_data = in_tensors_.at(SECOND_INPUT)->MutableData();
90 auto in_tensor_var_data = in_tensors_.at(THIRD_INPUT)->MutableData();
91 if (in_tensor_mean_data == nullptr || in_tensor_var_data == nullptr) {
92 FreeMeanAndVariance();
93 return RET_ERROR;
94 }
95 memcpy(mean_, in_tensor_mean_data, in_tensors_.at(SECOND_INPUT)->Size());
96 memcpy(variance_, in_tensor_var_data, in_tensors_.at(THIRD_INPUT)->Size());
97 return RET_OK;
98 }
99
Run()100 int BatchnormCPUKernel::Run() {
101 auto ret = ParallelLaunch(this->ms_context_, BatchNormRun, this, op_parameter_->thread_num_);
102 if (ret != RET_OK) {
103 MS_LOG(ERROR) << "BatchnormRun error error_code[" << ret << "]";
104 }
105 return ret;
106 }
107
DoExecute(int task_id)108 int BatchnormCPUKernel::DoExecute(int task_id) {
109 auto param = reinterpret_cast<BatchNormParameter *>(op_parameter_);
110 auto in_tensor_data = in_tensors_.at(0)->MutableData();
111 CHECK_NULL_RETURN(in_tensor_data);
112 auto out_tensor_data = out_tensors_.at(0)->MutableData();
113 CHECK_NULL_RETURN(out_tensor_data);
114 BatchNormFp32(in_tensor_data, mean_, variance_, param, task_id, out_tensor_data);
115 return RET_OK;
116 }
117
BatchNormRun(void * cdata,int task_id,float lhs_scale,float rhs_scale)118 int BatchNormRun(void *cdata, int task_id, float lhs_scale, float rhs_scale) {
119 CHECK_NULL_RETURN(cdata);
120 auto kernel = reinterpret_cast<BatchnormCPUKernel *>(cdata);
121 auto ret = kernel->DoExecute(task_id);
122 if (ret != RET_OK) {
123 MS_LOG(ERROR) << "BatchnormRun error task_id[" << task_id << "] error_code[" << ret << "]";
124 }
125 return ret;
126 }
127
set_momentum(float momentum)128 int BatchnormCPUKernel::set_momentum(float momentum) {
129 auto param = reinterpret_cast<BatchNormParameter *>(op_parameter_);
130 param->momentum_ = momentum;
131
132 return RET_OK;
133 }
134
get_momentum()135 float BatchnormCPUKernel::get_momentum() {
136 auto param = reinterpret_cast<BatchNormParameter *>(op_parameter_);
137 return param->momentum_;
138 }
139
RestoreDefaultMomentum()140 int BatchnormCPUKernel::RestoreDefaultMomentum() {
141 auto ret = set_momentum(default_momentum_);
142 if (ret != RET_OK) {
143 MS_LOG(ERROR) << "Restore Momentum Error";
144 return RET_ERROR;
145 }
146 return RET_OK;
147 }
148
149 REG_KERNEL(kCPU, kNumberTypeFloat32, PrimitiveType_BatchNorm, LiteKernelCreator<BatchnormCPUKernel>)
150 } // namespace mindspore::kernel
151