1. 自定义ReLU变体算子

任务目标

在TFLite Micro中集成一个自定义的Leaky ReLU算子（带斜率ReLU），并通过单元测试验证其输出正确性。本案例聚焦于算子内核实现、注册机制与测试闭环，不涉及训练或模型转换流程。

环境依赖

• TFLite Micro源码（含tensorflow/lite/micro目录）
• C++编译器（支持C++11及以上）
• Makefile构建系统（TFLite Micro默认使用）

实现方案

Leaky ReLU定义：y = x if x >= 0, else alpha * x。本实现将alpha作为算子内置参数（固定为0.1），避免引入额外tensor参数，保持内核最小化。注册为自定义算子kTfLiteCustom，kernel name为"LeakyRelu"。

关键完整代码

// leaky_relu_kernel.h
#ifndef TENSORFLOW_LITE_MICRO_EXAMPLES_LEAKY_RELU_KERNEL_H_
#define TENSORFLOW_LITE_MICRO_EXAMPLES_LEAKY_RELU_KERNEL_H_

#include "tensorflow/lite/c/common.h"
#include "tensorflow/lite/micro/micro_utils.h"

namespace tflite {

constexpr float kLeakyReluAlpha = 0.1f;

TfLiteStatus LeakyReluPrepare(TfLiteContext* context, TfLiteNode* node) {
  MicroContext* micro_context = GetMicroContext(context);
  TfLiteTensor* input = micro_context->AllocateTempInputTensor(node, 0);
  TfLiteTensor* output = micro_context->AllocateTempOutputTensor(node, 0);
  if (input == nullptr || output == nullptr) return kTfLiteError;
  if (input->type != kTfLiteFloat32 || output->type != kTfLiteFloat32) {
    return kTfLiteError;
  }
  output->params.scale = input->params.scale;
  output->params.zero_point = input->params.zero_point;
  micro_context->DeallocateTempTfLiteTensor(input);
  micro_context->DeallocateTempTfLiteTensor(output);
  return kTfLiteOk;
}

TfLiteStatus LeakyReluEval(TfLiteContext* context, TfLiteNode* node) {
  TfLiteEvalTensor* input = tflite::micro::GetEvalInput(context, node, 0);
  TfLiteEvalTensor* output = tflite::micro::GetEvalOutput(context, node, 0);
  const float* input_data = input->data.f;
  float* output_data = output->data.f;
  size_t num_elements = ElementCount(*input->dims);
  for (size_t i = 0; i < num_elements; ++i) {
    output_data[i] = input_data[i] >= 0.0f ? input_data[i] : kLeakyReluAlpha * input_data[i];
  }
  return kTfLiteOk;
}

}  // namespace tflite
#endif  // TENSORFLOW_LITE_MICRO_EXAMPLES_LEAKY_RELU_KERNEL_H_

// leaky_relu_test.cc
#include "tensorflow/lite/c/common.h"
#include "tensorflow/lite/micro/all_ops_resolver.h"
#include "tensorflow/lite/micro/micro_interpreter.h"
#include "tensorflow/lite/micro/testing/micro_test.h"
#include "tensorflow/lite/micro/examples/leaky_relu_kernel.h"

namespace tflite {
namespace {

TfLiteRegistration* Register_LeakyRelu() {
  static TfLiteRegistration r = {nullptr, nullptr, LeakyReluPrepare, LeakyReluEval};
  return &r;
}

class LeakyReluOpResolver : public AllOpsResolver {
 public:
  const TfLiteRegistration* FindOp(tflite::BuiltinOperator op,
                                   int version) const override {
    return AllOpsResolver::FindOp(op, version);
  }
  const TfLiteRegistration* FindOp(const char* custom_name,
                                   int version) const override {
    if (strcmp(custom_name, "LeakyRelu") == 0) return Register_LeakyRelu();
    return AllOpsResolver::FindOp(custom_name, version);
  }
};

}  // namespace
}  // namespace tflite

TF_LITE_MICRO_TESTS_BEGIN

TF_LITE_MICRO_TEST(LeakyReluPositiveInput) {
  float input_data[] = {0.0f, 1.0f, 2.5f};
  float output_data[3] = {0};
  TfLiteTensor tensors[] = {
      {nullptr, nullptr, kTfLiteFloat32, {3}, input_data, 0, {}},
      {nullptr, nullptr, kTfLiteFloat32, {3}, output_data, 0, {}},
  };
  TfLiteContext context{};
  TfLiteNode node{};
  node.inputs = {0};
  node.outputs = {1};
  node.custom_initial_data = nullptr;
  node.custom_initial_data_size = 0;
  node.user_data = nullptr;
  node.builtin_data = nullptr;
  node.delegate = nullptr;
  node.temporaries = {};

  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, LeakyReluPrepare(&context, &node));
  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, LeakyReluEval(&context, &node));
  TF_LITE_MICRO_EXPECT_EQ(0.0f, output_data[0]);
  TF_LITE_MICRO_EXPECT_EQ(1.0f, output_data[1]);
  TF_LITE_MICRO_EXPECT_EQ(2.5f, output_data[2]);
}

TF_LITE_MICRO_TEST(LeakyReluNegativeInput) {
  float input_data[] = {-1.0f, -0.5f, -2.0f};
  float output_data[3] = {0};
  TfLiteTensor tensors[] = {
      {nullptr, nullptr, kTfLiteFloat32, {3}, input_data, 0, {}},
      {nullptr, nullptr, kTfLiteFloat32, {3}, output_data, 0, {}},
  };
  TfLiteContext context{};
  TfLiteNode node{};
  node.inputs = {0};
  node.outputs = {1};

  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, LeakyReluPrepare(&context, &node));
  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, LeakyReluEval(&context, &node));
  TF_LITE_MICRO_EXPECT_EQ(-0.1f, output_data[0]);
  TF_LITE_MICRO_EXPECT_EQ(-0.05f, output_data[1]);
  TF_LITE_MICRO_EXPECT_EQ(-0.2f, output_data[2]);
}

TF_LITE_MICRO_TESTS_END

构建与运行

# 将上述两个文件放入 tensorflow/lite/micro/examples/ 目录
# 在TFLite Micro根目录执行：
make -f tensorflow/lite/micro/tools/make/Makefile \
  test_leaky_relu_test \
  TARGET=linux

验证方法

运行上述make命令后，测试框架自动执行两个测试用例。若输出包含All tests passed，则说明Leaky ReLU算子对正数输入保持原值、对负数输入乘以0.1的斜率，输出正确性得到验证。测试覆盖了边界值0、正数、负数三类典型输入。

常见问题

• 链接错误：确保leaky_relu_kernel.h与leaky_relu_test.cc位于同一目录，且Makefile能搜索到该路径。
• 算子未注册：若在完整推理中使用，需在AllOpsResolver子类中重载FindOp(const char*, int)，本测试已展示最小注册方式。
• 量化支持：本实现仅支持float32，如需int8量化需额外实现量化参数处理逻辑。