You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
167 lines
6.9 KiB
167 lines
6.9 KiB
4 months ago
|
/*
|
||
|
* Copyright (C) 2018 The Android Open Source Project
|
||
|
*
|
||
|
* 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.
|
||
|
*/
|
||
|
|
||
|
#define LOG_TAG "Operations"
|
||
|
|
||
|
#include <cmath>
|
||
|
#include <vector>
|
||
|
|
||
|
#include "OperationResolver.h"
|
||
|
#include "Tracing.h"
|
||
|
|
||
|
#ifdef NN_INCLUDE_CPU_IMPLEMENTATION
|
||
|
#include "CpuOperationUtils.h"
|
||
|
#endif // NN_INCLUDE_CPU_IMPLEMENTATION
|
||
|
|
||
|
namespace android {
|
||
|
namespace nn {
|
||
|
namespace instance_normalization {
|
||
|
|
||
|
constexpr char kOperationName[] = "INSTANCE_NORMALIZATION";
|
||
|
|
||
|
constexpr uint32_t kNumInputs = 5;
|
||
|
constexpr uint32_t kInputTensor = 0;
|
||
|
constexpr uint32_t kGammaScalar = 1;
|
||
|
constexpr uint32_t kBetaScalar = 2;
|
||
|
constexpr uint32_t kEpsilonScalar = 3;
|
||
|
constexpr uint32_t kLayoutScalar = 4;
|
||
|
|
||
|
constexpr uint32_t kNumOutputs = 1;
|
||
|
constexpr uint32_t kOutputTensor = 0;
|
||
|
|
||
|
#ifdef NN_INCLUDE_CPU_IMPLEMENTATION
|
||
|
namespace {
|
||
|
|
||
|
template <typename T>
|
||
|
inline bool instanceNormNhwc(const T* inputData, const Shape& inputShape, T gamma, T beta,
|
||
|
T epsilon, T* outputData, const Shape& outputShape) {
|
||
|
NNTRACE_TRANS("InstanceNormalizationNhwc");
|
||
|
uint32_t numBatches = getSizeOfDimension(inputShape, 0);
|
||
|
uint32_t height = getSizeOfDimension(inputShape, 1);
|
||
|
uint32_t width = getSizeOfDimension(inputShape, 2);
|
||
|
uint32_t depth = getSizeOfDimension(inputShape, 3);
|
||
|
for (uint32_t b = 0; b < numBatches; b++) {
|
||
|
for (uint32_t d = 0; d < depth; d++) {
|
||
|
uint32_t indexBase = b * height * width * depth + d;
|
||
|
T mean = 0, sigma = 0;
|
||
|
|
||
|
// Compute the mean of a single layer.
|
||
|
for (uint32_t h = 0; h < height; h++) {
|
||
|
for (uint32_t w = 0; w < width; w++) {
|
||
|
T val = inputData[indexBase + (h * width + w) * depth];
|
||
|
mean += val;
|
||
|
}
|
||
|
}
|
||
|
mean /= static_cast<T>(height * width);
|
||
|
|
||
|
// Compute the standard deviation (sigma) of a single layer.
|
||
|
for (uint32_t h = 0; h < height; h++) {
|
||
|
for (uint32_t w = 0; w < width; w++) {
|
||
|
T val = inputData[indexBase + (h * width + w) * depth] - mean;
|
||
|
sigma += val * val;
|
||
|
}
|
||
|
}
|
||
|
sigma = std::sqrt(static_cast<float>(sigma / static_cast<T>(height * width)) + epsilon);
|
||
|
|
||
|
// Apply instance normalization.
|
||
|
for (uint32_t h = 0; h < height; h++) {
|
||
|
for (uint32_t w = 0; w < width; w++) {
|
||
|
uint32_t ind = indexBase + (h * width + w) * depth;
|
||
|
outputData[ind] = (inputData[ind] - mean) * gamma / sigma + beta;
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
return true;
|
||
|
}
|
||
|
|
||
|
template <typename T>
|
||
|
inline bool instanceNorm(const T* inputData, const Shape& inputShape, T gamma, T beta, T epsilon,
|
||
|
bool useNchw, T* outputData, const Shape& outputShape) {
|
||
|
InputWithLayout<T> input(useNchw);
|
||
|
OutputWithLayout<T> output(useNchw);
|
||
|
NN_RET_CHECK(input.initialize(inputData, inputShape));
|
||
|
NN_RET_CHECK(output.initialize(outputData, outputShape));
|
||
|
NN_RET_CHECK(instanceNormNhwc(input.getNhwcBuffer(), input.getNhwcShape(), gamma, beta, epsilon,
|
||
|
output.getNhwcBuffer(), output.getNhwcShape()));
|
||
|
NN_RET_CHECK(output.commit());
|
||
|
return true;
|
||
|
}
|
||
|
|
||
|
} // namespace
|
||
|
#endif // NN_INCLUDE_CPU_IMPLEMENTATION
|
||
|
|
||
|
Result<Version> validate(const IOperationValidationContext* context) {
|
||
|
NN_RET_CHECK_EQ(context->getNumInputs(), kNumInputs);
|
||
|
NN_RET_CHECK_EQ(context->getNumOutputs(), kNumOutputs);
|
||
|
std::vector<OperandType> inExpectedTypes;
|
||
|
auto inputType = context->getInputType(kInputTensor);
|
||
|
if (inputType == OperandType::TENSOR_FLOAT32) {
|
||
|
inExpectedTypes = {OperandType::TENSOR_FLOAT32, OperandType::FLOAT32, OperandType::FLOAT32,
|
||
|
OperandType::FLOAT32, OperandType::BOOL};
|
||
|
} else if (inputType == OperandType::TENSOR_FLOAT16) {
|
||
|
inExpectedTypes = {OperandType::TENSOR_FLOAT16, OperandType::FLOAT16, OperandType::FLOAT16,
|
||
|
OperandType::FLOAT16, OperandType::BOOL};
|
||
|
} else {
|
||
|
return NN_ERROR() << "Unsupported input tensor type for operation " << kOperationName;
|
||
|
}
|
||
|
NN_RET_CHECK(validateInputTypes(context, inExpectedTypes));
|
||
|
NN_RET_CHECK(validateOutputTypes(context, {inputType}));
|
||
|
return Version::ANDROID_Q;
|
||
|
}
|
||
|
|
||
|
#ifdef NN_INCLUDE_CPU_IMPLEMENTATION
|
||
|
bool prepare(IOperationExecutionContext* context) {
|
||
|
Shape input = context->getInputShape(kInputTensor);
|
||
|
NN_RET_CHECK_EQ(getNumberOfDimensions(input), 4);
|
||
|
return context->setOutputShape(kOutputTensor, input);
|
||
|
}
|
||
|
|
||
|
bool execute(IOperationExecutionContext* context) {
|
||
|
switch (context->getInputType(kInputTensor)) {
|
||
|
case OperandType::TENSOR_FLOAT16:
|
||
|
return instanceNorm(context->getInputBuffer<_Float16>(kInputTensor),
|
||
|
context->getInputShape(kInputTensor),
|
||
|
context->getInputValue<_Float16>(kGammaScalar),
|
||
|
context->getInputValue<_Float16>(kBetaScalar),
|
||
|
context->getInputValue<_Float16>(kEpsilonScalar),
|
||
|
context->getInputValue<bool>(kLayoutScalar),
|
||
|
context->getOutputBuffer<_Float16>(kOutputTensor),
|
||
|
context->getOutputShape(kOutputTensor));
|
||
|
case OperandType::TENSOR_FLOAT32:
|
||
|
return instanceNorm(context->getInputBuffer<float>(kInputTensor),
|
||
|
context->getInputShape(kInputTensor),
|
||
|
context->getInputValue<float>(kGammaScalar),
|
||
|
context->getInputValue<float>(kBetaScalar),
|
||
|
context->getInputValue<float>(kEpsilonScalar),
|
||
|
context->getInputValue<bool>(kLayoutScalar),
|
||
|
context->getOutputBuffer<float>(kOutputTensor),
|
||
|
context->getOutputShape(kOutputTensor));
|
||
|
default:
|
||
|
NN_RET_CHECK_FAIL() << "Unsupported tensor type for operation " << kOperationName;
|
||
|
}
|
||
|
}
|
||
|
#endif // NN_INCLUDE_CPU_IMPLEMENTATION
|
||
|
|
||
|
} // namespace instance_normalization
|
||
|
|
||
|
NN_REGISTER_OPERATION(INSTANCE_NORMALIZATION, instance_normalization::kOperationName,
|
||
|
instance_normalization::validate, instance_normalization::prepare,
|
||
|
instance_normalization::execute);
|
||
|
|
||
|
} // namespace nn
|
||
|
} // namespace android
|