如何在TensorFlow Lite Micro中添加自定义操作符（2）

上一篇中，小编给大家抽丝剥茧的介绍了在TFLm中实现一个算子所涉及的文件，以及每个文件的具体作用，包括：功能实现，算子解析等。那么本篇就带着大家一起看下注册机制是怎么实现的？我们还是先以reshape算子进行说明，如何将reshape算子注册到解析器中，接下来介绍如果我们想自定义一个算子需要干些什么。

  操作符注册到解析器

1.1 在 MicroMutableOpResolver 中添加注册方法

文件位置：`micro/micro_mutable_op_resolver.h`，在类定义中添加以下方法：

 

TfLiteStatus AddReshape() {
  return AddBuiltin(BuiltinOperator_RESHAPE,
                    tflite::Register_RESHAPE(), ParseReshape);
}

 

注册方法说明：

AddBuiltin 函数：MicroMutableOpResolver 的核心方法，用于注册内置操作符

BuiltinOperator_RESHAPE：操作符的唯一标识符，与 FlatBuffer schema 中的定义一致

Register_RESHAPE()：返回操作符的注册信息，包含执行函数指针

ParseReshape：参数解析函数指针，用于从模型文件中解析参数

1.2 在全局解析器中注册

文件位置：`micro/all_ops_resolver.cpp`，在 `AllOpsResolver` 构造函数中添加：

 

AddReshape();

 

全局注册说明：

`AllOpsResolver` 包含了所有标准 TFLite 操作符

适用于需要完整操作符支持的应用场景

会增加代码大小，但提供最大的模型兼容性

  添加新操作符的完整步骤

步骤 1：创建内核实现文件

创建文件：`micro/kernels/your_op.cpp`

1. 包含必要的头文件：

 

 #include "tensorflow/lite/c/builtin_op_data.h"
   #include "tensorflow/lite/c/common.h"
   #include "tensorflow/lite/micro/kernels/kernel_util.h"
   // 其他必要的头文件

 

2. 定义命名空间和常量：

 

 namespace tflite {
   namespace ops {
   namespace micro {
   namespace your_op {
   constexpr int kInputTensor = 0;
   constexpr int kOutputTensor = 0;
   // 其他常量定义

 

3. 实现核心函数：

`Prepare` 函数：验证参数，计算输出形状

`Eval` 函数：执行实际计算

可选的 `Init` 函数：如果需要持久化数据

4. 创建注册函数：

 

TfLiteRegistration_V1 Register_YOUR_OP() {
     return tflite::RegisterOp(Init, Prepare, Eval);
   }

 

步骤2：注册操作符

修改文件：`micro/micro_mutable_op_resolver.h`

在类定义中添加注册方法：

 

TfLiteStatus AddYourOp() {
  return AddBuiltin(BuiltinOperator_YOUR_OP,
                    tflite::Register_YOUR_OP(), ParseYourOp);
}

 

修改文件：`micro/all_ops_resolver.cpp`

在构造函数中添加：

 

AddYourOp();

  关键注意事项

 

内存管理最佳实践

1. 临时张量管理：

 

 // 正确的临时张量使用方式
   TfLiteTensor* input = micro_context->AllocateTempInputTensor(node, 0);
   // 使用张量...
   micro_context->DeallocateTempTfLiteTensor(input);  // 必须释放

 

2. 持久化内存 vs 临时内存：

持久化内存：用于存储操作符参数、权重等需要长期保存的数据

临时内存：用于计算过程中的中间结果，使用后立即释放

3. 内存对齐：

微控制器对内存对齐有严格要求

使用 `MicroArenaBufferAlignment()` 获取正确的对齐值

  错误处理规范

1. 参数验证：

 

TF_LITE_ENSURE(context, condition);           // 条件检查
   TF_LITE_ENSURE_EQ(context, actual, expected); // 相等性检查
   TF_LITE_ENSURE_STATUS(status);                // 状态码检查

 

2. 错误报告：

 

TF_LITE_KERNEL_LOG(context, "Error message with details");

 

3. 状态码使用：

`kTfLiteOk`：操作成功

`kTfLiteError`：一般错误

`kTfLiteDelegateError`：委托相关错误

  性能优化策略

1. 避免重复计算：

在 Prepare 阶段完成形状计算

缓存经常使用的计算结果

2. 内存访问优化：

尽量使用连续内存访问模式

避免频繁的小块内存分配

3. 原地操作：

当可能时，使用原地操作减少内存拷贝

检查输入输出是否可以共享内存

4. 数据类型优化：

支持量化数据类型（int8, uint8）

针对不同数据类型提供优化实现

通过遵循以上流程，我们就可以是现在 TensorFlow Lite Micro中添加自定义操作符的操作了，并确保其在资源受限的微控制器环境中稳定高效地运行。

这样一来，就可以不被TFLm的原生算子的约束，放开手脚运行更好的模型。一起探讨，让我们更懂TFLm，更懂模型！！！