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

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

操作符注册到解析器

1.1 在 MicroMutableOpResolver 中添加注册方法

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

TfLiteStatusAddReshape() {
returnAddBuiltin(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. 定义命名空间和常量：

namespacetflite {
 namespaceops {
 namespacemicro {
 namespaceyour_op {
 constexprintkInputTensor =0;
 constexprintkOutputTensor =0;
 // 其他常量定义

3. 实现核心函数：

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

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

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

4. 创建注册函数：

TfLiteRegistration_V1Register_YOUR_OP() {
  returntflite::RegisterOp(Init,Prepare,Eval);
 }

步骤2：注册操作符

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

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

TfLiteStatusAddYourOp() {
returnAddBuiltin(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，更懂模型！！！