TensorFlow Lite Micro中为算子传递自定义数据的实现方案

在嵌入式AI应用中，TensorFlow Lite Micro(TFLM)作为轻量级推理框架被广泛使用。本文探讨了如何在TFLM中为内置算子传递额外参数的技术实现方案，这对于专用计算模块的集成尤为重要。

问题背景

在神经网络模型部署过程中，某些层(如线性层)可能通过专用计算模块获得优化。这些计算模块往往需要额外的配置信息，例如：

• 计算核心的分配策略
• 数据分块方案
• 硬件特定参数

传统做法是创建自定义算子(Custom OP)，但这会带来模型兼容性问题，且需要重新训练模型。

技术方案比较

方案一：自定义算子

1. 为每个可优化层创建自定义算子
2. 替换原始层并重新训练模型
3. 转换模型时添加自定义算子及参数
4. 实现对应的算子逻辑

缺点：破坏模型兼容性，需要重新训练。

方案二：元数据扩展

利用TFLite模型的Metadata机制，这是一种键值存储结构，可以附加任意数据。优势在于：

• 不修改原始算子结构
• 参考实现可忽略元数据
• 支持更灵活的数据组织方式

关键技术实现

核心挑战是如何在算子执行时获取对应的元数据。解决方案包括：

1. 上下文扩展：通过MicroContext提供当前执行的子图和算子索引
2. 索引追踪：维护current_subgraph_idx和previous_subgraph_idx状态
3. 元数据查询：基于子图和算子索引定位对应的配置数据

实现细节

在TFLM中，算子可以通过以下方式获取上下文信息：

tflite::GetMicroContext(context)->graph()->GetCurrentSubgraphIndex();
tflite::GetMicroContext(context)->graph()->GetCurrentOperatorIndex(); 

专用计算模块实现者可以：

1. 在模型转换阶段通过AOT工具注入元数据
2. 在算子实现中查询并使用这些参数
3. 保持与参考实现的兼容性

应用价值

该方案特别适合以下场景：

• 异构计算平台部署
• 专用计算模块集成
• 动态计算资源配置
• 性能调优实验

相比自定义算子方案，元数据扩展保持了模型兼容性，降低了部署复杂度，是更优雅的解决方案。