Energy op

python/tflite_micro/python_ops_resolver.cc

@@ -45,6 +45,7 @@ PythonOpsResolver::PythonOpsResolver() {
  AddDequantize();
  AddDetectionPostprocess();
  AddDiv();
+ AddEnergy();
  AddElu();
  AddEqual();
  AddEthosU();

signal/micro/kernels/BUILD

@@ -9,6 +9,7 @@ cc_library(
  name = "register_signal_ops",
  srcs = [
  "delay.cc",
+ "energy.cc",
  "framer.cc",
  "overlap_add.cc",
  "rfft.cc",
@@ -24,6 +25,7 @@ cc_library(
  ],
  deps = [
  "//signal/src:circular_buffer",
+ "//signal/src:energy",
  "//signal/src:overlap_add",
  "//signal/src:rfft",
  "//signal/src:window",
@@ -194,3 +196,29 @@ cc_test(
  "//tensorflow/lite/micro/testing:micro_test",
  ],
 )
+
+cc_library(
+ name = "energy_flexbuffers_generated_data",
+ srcs = [
+ "energy_flexbuffers_generated_data.cc",
+ ],
+ hdrs = [
+ "energy_flexbuffers_generated_data.h",
+ ],
+)
+
+cc_test(
+ name = "energy_test",
+ srcs = [
+ "energy_test.cc",
+ ],
+ deps = [
+ ":energy_flexbuffers_generated_data",
+ ":register_signal_ops",
+ "//tensorflow/lite/c:common",
+ "//tensorflow/lite/micro:op_resolvers",
+ "//tensorflow/lite/micro:test_helpers",
+ "//tensorflow/lite/micro/kernels:kernel_runner",
+ "//tensorflow/lite/micro/testing:micro_test",
+ ],
+)

signal/micro/kernels/energy.cc

@@ -0,0 +1,112 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+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 "signal/src/energy.h"
+
+#include <math.h>
+#include <stddef.h>
+#include <stdint.h>
+
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/micro/flatbuffer_utils.h"
+#include "tensorflow/lite/micro/kernels/kernel_util.h"
+#include "tensorflow/lite/micro/micro_context.h"
+
+namespace tflite {
+namespace {
+
+constexpr int kInputTensor = 0;
+constexpr int kOutputTensor = 0;
+
+// Indices into the init flexbuffer's vector.
+// The parameter's name is in the comment that follows.
+// Elements in the vectors are ordered alphabetically by parameter name.
+constexpr int kEndIndexIndex = 0; // 'end_index'
+constexpr int kStartIndexIndex = 1; // 'start_index'
+
+struct TFLMSignalEnergyParams {
+ int32_t end_index;
+ int32_t start_index;
+};
+
+void* Init(TfLiteContext* context, const char* buffer, size_t length) {
+ TFLITE_DCHECK(context->AllocatePersistentBuffer != nullptr);
+
+ auto* data =
+ static_cast<TFLMSignalEnergyParams*>(context->AllocatePersistentBuffer(
+ context, sizeof(TFLMSignalEnergyParams)));
+
+ if (data == nullptr) {
+ return nullptr;
+ }
+
+ tflite::FlexbufferWrapper fbw(reinterpret_cast<const uint8_t*>(buffer),
+ length);
+ data->end_index = fbw.ElementAsInt32(kEndIndexIndex);
+ data->start_index = fbw.ElementAsInt32(kStartIndexIndex);
+ return data;
+}
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+ TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
+ TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+
+ MicroContext* micro_context = GetMicroContext(context);
+ TfLiteTensor* input =
+ micro_context->AllocateTempInputTensor(node, kInputTensor);
+ TF_LITE_ENSURE(context, input != nullptr);
+ TfLiteTensor* output =
+ micro_context->AllocateTempOutputTensor(node, kOutputTensor);
+ TF_LITE_ENSURE(context, output != nullptr);
+
+ TF_LITE_ENSURE_EQ(context, NumDimensions(input), 1);
+ TF_LITE_ENSURE_EQ(context, NumDimensions(output), 1);
+
+ TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteInt16);
+ TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteUInt32);
+
+ micro_context->DeallocateTempTfLiteTensor(input);
+ micro_context->DeallocateTempTfLiteTensor(output);
+ return kTfLiteOk;
+}
+
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+ auto* params = reinterpret_cast<TFLMSignalEnergyParams*>(node->user_data);
+
+ const TfLiteEvalTensor* input =
+ tflite::micro::GetEvalInput(context, node, kInputTensor);
+ TfLiteEvalTensor* output =
+ tflite::micro::GetEvalOutput(context, node, kOutputTensor);
+
+ const Complex<int16_t>* input_data =
+ tflite::micro::GetTensorData<Complex<int16_t>>(input);
+ uint32_t* output_data = tflite::micro::GetTensorData<uint32_t>(output);
+
+ tflm_signal::SpectrumToEnergy(input_data, params->start_index,
+ params->end_index, output_data);
+ return kTfLiteOk;
+}
+
+} // namespace
+
+namespace tflm_signal {
+TFLMRegistration* Register_ENERGY() {
+ static TFLMRegistration r = tflite::micro::RegisterOp(Init, Prepare, Eval);
+ return &r;
+}
+} // namespace tflm_signal
+
+} // namespace tflite

signal/micro/kernels/energy_flexbuffers_generated_data.cc

@@ -0,0 +1,37 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+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.
+==============================================================================*/
+
+// This file is generated. See:
+// tensorflow/lite/micro/kernels/test_data_generation/README.md
+
+#include "signal/micro/kernels/energy_flexbuffers_generated_data.h"
+
+const int g_gen_data_size_start_index_2_end_index_4 = 35;
+const unsigned char g_gen_data_start_index_2_end_index_4[] = {
+ 0x73, 0x74, 0x61, 0x72, 0x74, 0x5f, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x00,
+ 0x65, 0x6e, 0x64, 0x5f, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x00, 0x02, 0x0b,
+ 0x18, 0x02, 0x01, 0x02, 0x04, 0x02, 0x04, 0x04, 0x04, 0x24, 0x01,
+};
+const int g_gen_data_size_start_index_0_end_index_4 = 35;
+const unsigned char g_gen_data_start_index_0_end_index_4[] = {
+ 0x73, 0x74, 0x61, 0x72, 0x74, 0x5f, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x00,
+ 0x65, 0x6e, 0x64, 0x5f, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x00, 0x02, 0x0b,
+ 0x18, 0x02, 0x01, 0x02, 0x04, 0x00, 0x04, 0x04, 0x04, 0x24, 0x01,
+};
+const int g_gen_data_size_start_index_4_end_index_8 = 35;
+const unsigned char g_gen_data_start_index_4_end_index_8[] = {
+ 0x73, 0x74, 0x61, 0x72, 0x74, 0x5f, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x00,
+ 0x65, 0x6e, 0x64, 0x5f, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x00, 0x02, 0x0b,
+ 0x18, 0x02, 0x01, 0x02, 0x08, 0x04, 0x04, 0x04, 0x04, 0x24, 0x01,
+};

signal/micro/kernels/energy_flexbuffers_generated_data.h

@@ -0,0 +1,28 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+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 SIGNAL_MICRO_KERNELS_TEST_DATA_GENERATION_GENERATE_ENERGY_FLEXBUFFERS_DATA_H_
+#define SIGNAL_MICRO_KERNELS_TEST_DATA_GENERATION_GENERATE_ENERGY_FLEXBUFFERS_DATA_H_
+
+extern const int g_gen_data_size_start_index_2_end_index_4;
+extern const unsigned char g_gen_data_start_index_2_end_index_4[];
+
+extern const int g_gen_data_size_start_index_0_end_index_4;
+extern const unsigned char g_gen_data_start_index_0_end_index_4[];
+
+extern const int g_gen_data_size_start_index_4_end_index_8;
+extern const unsigned char g_gen_data_start_index_4_end_index_8[];
+
+#endif // SIGNAL_MICRO_KERNELS_TEST_DATA_GENERATION_GENERATE_ENERGY_FLEXBUFFERS_DATA_H_

signal/micro/kernels/energy_test.cc

@@ -0,0 +1,123 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+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 "signal/micro/kernels/energy_flexbuffers_generated_data.h"
+#include "tensorflow/lite/micro/kernels/kernel_runner.h"
+#include "tensorflow/lite/micro/test_helpers.h"
+#include "tensorflow/lite/micro/testing/micro_test.h"
+
+namespace tflite {
+namespace testing {
+namespace {
+
+TfLiteStatus TestEnergy(int* input_dims_data, const int16_t* input_data,
+ int* output_dims_data, const uint32_t* golden,
+ const unsigned char* flexbuffers_data,
+ const unsigned int flexbuffers_data_size,
+ uint32_t* output_data) {
+ TfLiteIntArray* input_dims = IntArrayFromInts(input_dims_data);
+ TfLiteIntArray* output_dims = IntArrayFromInts(output_dims_data);
+ const int output_len = ElementCount(*output_dims);
+ constexpr int kInputsSize = 1;
+ constexpr int kOutputsSize = 1;
+ constexpr int kTensorsSize = kInputsSize + kOutputsSize;
+ TfLiteTensor tensors[kTensorsSize] = {
+ CreateTensor(input_data, input_dims),
+ CreateTensor(output_data, output_dims),
+ };
+
+ int inputs_array_data[] = {1, 0};
+ TfLiteIntArray* inputs_array = IntArrayFromInts(inputs_array_data);
+ int outputs_array_data[] = {1, 1};
+ TfLiteIntArray* outputs_array = IntArrayFromInts(outputs_array_data);
+
+ const TFLMRegistration* registration = tflm_signal::Register_ENERGY();
+ micro::KernelRunner runner(*registration, tensors, kTensorsSize, inputs_array,
+ outputs_array,
+ /*builtin_data=*/nullptr);
+
+ // TfLite uses a char* for the raw bytes whereas flexbuffers use an unsigned
+ // char*. This small discrepancy results in compiler warnings unless we
+ // reinterpret_cast right before passing in the flexbuffer bytes to the
+ // KernelRunner.
+ TfLiteStatus status = runner.InitAndPrepare(
+ reinterpret_cast<const char*>(flexbuffers_data), flexbuffers_data_size);
+ if (status != kTfLiteOk) {
+ return status;
+ }
+
+ status = runner.Invoke();
+ if (status != kTfLiteOk) {
+ return status;
+ }
+
+ for (int i = 0; i < output_len; ++i) {
+ TF_LITE_MICRO_EXPECT_EQ(golden[i], output_data[i]);
+ }
+ return kTfLiteOk;
+}
+
+} // namespace
+} // namespace testing
+} // namespace tflite
+
+TF_LITE_MICRO_TESTS_BEGIN
+
+TF_LITE_MICRO_TEST(EnergyTestMiddle) {
+ int input_shape[] = {1, 16};
+ int output_shape[] = {1, 8};
+ const int16_t input[] = {1, 2, 3, 4, 5, 6, 7, 8,
+ 9, 10, 11, 12, 13, 14, 15, 16};
+ const uint32_t golden[] = {0, 0, 61, 113, 0, 0, 0, 0};
+ uint32_t output[8];
+ memset(output, 0, sizeof(output));
+ TF_LITE_MICRO_EXPECT_EQ(
+ kTfLiteOk, tflite::testing::TestEnergy(
+ input_shape, input, output_shape, golden,
+ g_gen_data_start_index_2_end_index_4,
+ g_gen_data_size_start_index_2_end_index_4, output));
+}
+
+TF_LITE_MICRO_TEST(EnergyTestStart) {
+ int input_shape[] = {1, 16};
+ int output_shape[] = {1, 8};
+ const int16_t input[] = {1, 2, 3, 4, 5, 6, 7, 8,
+ 9, 10, 11, 12, 13, 14, 15, 16};
+ const uint32_t golden[] = {5, 25, 61, 113, 0, 0, 0, 0};
+ uint32_t output[8];
+ memset(output, 0, sizeof(output));
+ TF_LITE_MICRO_EXPECT_EQ(
+ kTfLiteOk, tflite::testing::TestEnergy(
+ input_shape, input, output_shape, golden,
+ g_gen_data_start_index_0_end_index_4,
+ g_gen_data_size_start_index_0_end_index_4, output));
+}
+
+TF_LITE_MICRO_TEST(EnergyTestEnd) {
+ int input_shape[] = {1, 16};
+ int output_shape[] = {1, 8};
+ const int16_t input[] = {1, 2, 3, 4, 5, 6, 7, 8,
+ 9, 10, 11, 12, 13, 14, 15, 16};
+ const uint32_t golden[] = {0, 0, 0, 0, 181, 265, 365, 481};
+ uint32_t output[8];
+ memset(output, 0, sizeof(output));
+ TF_LITE_MICRO_EXPECT_EQ(
+ kTfLiteOk, tflite::testing::TestEnergy(
+ input_shape, input, output_shape, golden,
+ g_gen_data_start_index_4_end_index_8,
+ g_gen_data_size_start_index_4_end_index_8, output));
+}
+
+TF_LITE_MICRO_TESTS_END

signal/src/BUILD

@@ -39,3 +39,10 @@ cc_library(
  srcs = ["overlap_add.cc"],
  hdrs = ["overlap_add.h"],
 )
+
+cc_library(
+ name = "energy",
+ srcs = ["energy.cc"],
+ hdrs = ["energy.h"],
+ deps = [":complex"],
+)