1 //
2 // CPUInstanceNorm.hpp
3 // MNN
4 //
5 // Created by MNN on 2019/02/28.
6 // Copyright © 2018, Alibaba Group Holding Limited
7 //
8
9 #include "backend/cpu/CPUInstanceNorm.hpp"
10 #include <math.h>
11 #include "backend/cpu/CPUBackend.hpp"
12 #include "core/Concurrency.h"
13 #include <MNN/MNNDefine.h>
14 #include "core/Macro.h"
15 #include "core/TensorUtils.hpp"
16
17 #ifdef MNN_USE_NEON
18 #include <arm_neon.h>
19 #endif
20
21 namespace MNN {
22
CPUInstanceNorm(Backend * backend,const MNN::Op * op)23 CPUInstanceNorm::CPUInstanceNorm(Backend* backend, const MNN::Op* op) : Execution(backend) {
24 auto normParam = op->main_as_BatchNorm();
25 const int channels = normParam->channels();
26 mEpsilon = normParam->epsilon();
27 mScale.reset(ALIGN_UP4(channels));
28 mScale.clear();
29 if (normParam->slopeData() && normParam->slopeData()->data()) {
30 ::memcpy(mScale.get(), normParam->slopeData()->data(), channels * sizeof(float));
31 }
32
33 mBias.reset(ALIGN_UP4(channels));
34 mBias.clear();
35 if (normParam->biasData() && normParam->biasData()->data()) {
36 ::memcpy(mBias.get(), normParam->biasData()->data(), channels * sizeof(float));
37 }
38 }
39
onExecute(const std::vector<Tensor * > & inputs,const std::vector<Tensor * > & outputs)40 ErrorCode CPUInstanceNorm::onExecute(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
41 MNN_ASSERT(3 == inputs.size());
42 MNN_ASSERT(1 == outputs.size());
43 auto input = inputs[0];
44 MNN_ASSERT(MNN_DATA_FORMAT_NC4HW4 == TensorUtils::getDescribe(input)->dimensionFormat);
45 auto mean = inputs[1];
46 auto variance = inputs[2];
47 auto output = outputs[0];
48 const int batch = input->batch();
49 const int batchStride = input->stride(0);
50 const int channelsDiv4 = UP_DIV(input->channel(), 4);
51 const int inImageSize = input->stride(1);
52 const float* scalePtr = mScale.get();
53 const float* biasPtr = mBias.get();
54 const float* meanPtr = mean->host<float>();
55 const float* variancePtr = variance->host<float>();
56
57 for (int b = 0; b < batch; ++b) {
58 const float* batchMeanPtr = meanPtr + b * mean->stride(0);
59 const float* batchVariancePtr = variancePtr + b * variance->stride(0);
60 const float* batchInputPtr = input->host<float>() + b * batchStride;
61 float* batchOutputPtr = output->host<float>() + b * batchStride;
62 MNN_CONCURRENCY_BEGIN(ic, channelsDiv4);
63 const int channelOffset = (int)ic << 2;
64 const float* channelsInputPtr = batchInputPtr + channelOffset * inImageSize;
65 float* channelsOutputPtr = batchOutputPtr + channelOffset * inImageSize;
66 #ifdef MNN_USE_NEON
67 float32x4_t epsilon = vdupq_n_f32(mEpsilon);
68 float32x4_t batchVariance = vld1q_f32(batchVariancePtr + channelOffset);
69 float32x4_t meanValue = vld1q_f32(batchMeanPtr + channelOffset);
70 float32x4_t scaleValue = vld1q_f32(scalePtr + channelOffset);
71 float32x4_t biasVaule = vld1q_f32(biasPtr + channelOffset);
72 float32x4_t rsqrt = vrsqrteq_f32(batchVariance + epsilon);
73
74 float32x4_t gamma = vmulq_f32(scaleValue, rsqrt);
75 float32x4_t beta = vsubq_f32(biasVaule, meanValue * gamma);
76 for (int i = 0; i < inImageSize; ++i) {
77 float32x4_t value = vld1q_f32(channelsInputPtr + i * 4);
78 vst1q_f32(channelsOutputPtr + i * 4, value * gamma + beta);
79 }
80
81 #else
82 float gamma[4];
83 float beta[4];
84 for (int k = 0; k < 4; ++k) {
85 const int index = channelOffset + k;
86 gamma[k] = scalePtr[index] / sqrt(batchVariancePtr[index] + mEpsilon);
87 beta[k] = biasPtr[index] - scalePtr[index] * batchMeanPtr[index] / sqrt(batchVariancePtr[index] + mEpsilon);
88 }
89
90 for (int i = 0; i < inImageSize; ++i) {
91 for (int k = 0; k < 4; ++k) {
92 channelsOutputPtr[i * 4 + k] = channelsInputPtr[i * 4 + k] * gamma[k] + beta[k];
93 }
94 }
95 #endif
96 MNN_CONCURRENCY_END();
97 }
98
99 return NO_ERROR;
100 }
101
102 class CPUInstanceNormCreator : public CPUBackend::Creator {
103 public:
onCreate(const std::vector<Tensor * > & inputs,const std::vector<Tensor * > & outputs,const MNN::Op * op,Backend * backend) const104 virtual Execution* onCreate(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
105 const MNN::Op* op, Backend* backend) const override {
106 return new CPUInstanceNorm(backend, op);
107 }
108 };
109
110 REGISTER_CPU_OP_CREATOR(CPUInstanceNormCreator, OpType_InstanceNorm);
111
112 } // namespace MNN
113