4. 自定义Mish激活算子

任务目标

在TFLite Micro中注册并实现自定义Mish激活算子，使其能够被推理引擎正确调用，并通过与标准Mish公式的逐元素对比验证输出一致性。

环境依赖

• TFLite Micro源码（含tensorflow/lite/micro目录）
• 编译器：arm-none-eabi-gcc 或 gcc（取决于目标平台）
• 测试框架：TFLite Micro自带的unit_test框架

实现步骤

Mish激活函数定义为：f(x) = x * tanh(softplus(x))，其中softplus(x) = ln(1 + e^x)。为避免数值溢出，当x > 20时softplus近似为x，此时tanh趋近于1，Mish近似为x。

关键完整代码

// mish_ops.h
#ifndef TENSORFLOW_LITE_MICRO_EXAMPLES_MISH_OPS_H_
#define TENSORFLOW_LITE_MICRO_EXAMPLES_MISH_OPS_H_

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

namespace tflite {
namespace ops {
namespace micro {
namespace custom {

TfLiteRegistration* Register_MISH();

}  // namespace custom
}  // namespace micro
}  // namespace ops
}  // namespace tflite

#endif  // TENSORFLOW_LITE_MICRO_EXAMPLES_MISH_OPS_H_

// mish_ops.cc
#include "tensorflow/lite/micro/examples/mish_ops.h"
#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
#include "tensorflow/lite/kernels/kernel_util.h"
#include "tensorflow/lite/micro/kernels/kernel_util.h"
#include <cmath>

namespace tflite {
namespace ops {
namespace micro {
namespace custom {
namespace mish {

constexpr int kInputTensor = 0;
constexpr int kOutputTensor = 0;

TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
  TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
  const TfLiteTensor* input = GetInput(context, node, kInputTensor);
  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
  TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteFloat32);
  TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteFloat32);
  output->dims->size = input->dims->size;
  for (int i = 0; i < input->dims->size; ++i) {
    output->dims->data[i] = input->dims->data[i];
  }
  return kTfLiteOk;
}

TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
  const TfLiteEvalTensor* input = tflite::micro::GetEvalInput(context, node, kInputTensor);
  TfLiteEvalTensor* output = tflite::micro::GetEvalOutput(context, node, kOutputTensor);
  const float* input_data = input->data.f;
  float* output_data = output->data.f;
  int num_elements = 1;
  for (int i = 0; i < input->dims->size; ++i) {
    num_elements *= input->dims->data[i];
  }
  for (int i = 0; i < num_elements; ++i) {
    float x = input_data[i];
    float softplus = (x > 20.0f) ? x : std::log1p(std::exp(x));
    float mish = x * std::tanh(softplus);
    output_data[i] = mish;
  }
  return kTfLiteOk;
}

}  // namespace mish

TfLiteRegistration* Register_MISH() {
  static TfLiteRegistration r = {nullptr, nullptr, mish::Prepare, mish::Eval};
  return &r;
}

}  // namespace custom
}  // namespace micro
}  // namespace ops
}  // namespace tflite

注册算子到解释器

// 在micro_interpreter.cc或模型解析注册表中添加
// 假设已有MicroMutableOpResolver实例resolver
resolver.AddCustom("Mish", tflite::ops::micro::custom::Register_MISH());

测试验证

// mish_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/mish_ops.h"
#include <cmath>

TF_LITE_MICRO_TESTS_BEGIN

TF_LITE_MICRO_TEST(MishMatchesReference) {
  // 输入测试数据：-5, -1, 0, 1, 5, 20
  float input_data[] = {-5.0f, -1.0f, 0.0f, 1.0f, 5.0f, 20.0f};
  float expected_output[6];
  for (int i = 0; i < 6; ++i) {
    float x = input_data[i];
    float sp = (x > 20.0f) ? x : std::log1p(std::exp(x));
    expected_output[i] = x * std::tanh(sp);
  }

  // 构建模型（单输入单输出，无权重）
  constexpr int kTensorArenaSize = 1024;
  uint8_t tensor_arena[kTensorArenaSize];
  tflite::MicroMutableOpResolver<1> resolver;
  resolver.AddCustom("Mish", tflite::ops::micro::custom::Register_MISH());

  // 此处需加载包含Mish节点的FlatBuffer模型
  // 为简化测试，直接构造调用Eval的测试（省略模型构建细节）
  // 实际测试应通过MicroInterpreter运行完整模型
  // 以下为直接调用Eval的单元测试方式：
  TfLiteTensor input_tensor, output_tensor;
  input_tensor.type = kTfLiteFloat32;
  input_tensor.dims = TfLiteIntArrayCreate(1);
  input_tensor.dims->data[0] = 6;
  input_tensor.data.f = input_data;
  output_tensor.type = kTfLiteFloat32;
  output_tensor.dims = TfLiteIntArrayCreate(1);
  output_tensor.dims->data[0] = 6;
  float output_data[6];
  output_tensor.data.f = output_data;

  TfLiteNode node;
  node.inputs = TfLiteIntArrayCreate(1);
  node.inputs->data[0] = 0;
  node.outputs = TfLiteIntArrayCreate(1);
  node.outputs->data[0] = 0;
  TfLiteContext context;
  context.tensors = new TfLiteTensor[2];
  context.tensors[0] = input_tensor;
  context.tensors[1] = output_tensor;

  auto* reg = tflite::ops::micro::custom::Register_MISH();
  reg->prepare(&context, &node);
  reg->eval(&context, &node);

  for (int i = 0; i < 6; ++i) {
    TF_LITE_MICRO_EXPECT_NEAR(expected_output[i], output_data[i], 1e-5f);
  }
  delete[] context.tensors;
}

TF_LITE_MICRO_TESTS_END

运行与验证

1. 将mish_ops.h和mish_ops.cc放入tensorflow/lite/micro/examples/目录。
2. 在BUILD文件中添加对应构建目标，或使用Makefile编译测试：
   make -f tensorflow/lite/micro/tools/make/Makefile test_mish_test
3. 运行编译产物，观察输出：
   ./gen/linux_x86_64_default/bin/mish_test
4. 通过条件：所有TF_LITE_MICRO_EXPECT_NEAR断言通过，即输出与标准Mish公式误差小于1e-5。

常见问题

• 数值溢出：当x较大时exp(x)溢出，代码中已用x > 20分支规避。
• 算子未注册：确保在解释器初始化时调用resolver.AddCustom("Mish", ...)，且模型中的算子名与注册名一致。
• 维度不匹配：Prepare中已复制输入维度到输出，若模型输入为动态形状需额外处理。