相信大家在部署嵌入式端的AI應用時，一定使用過TensorFlow Lite Micro，以下簡稱TFLm。TFLm 是專為微控制器和嵌入式設備設計的輕量級機器學習推理框架，它通過模塊化的操作符系統來支持各種神經網絡層的計算。也就是說，我們不僅可以使用內嵌的算子運算，還可以自己注冊一個新的算子，更加的靈活。本期就將用兩期的文章以 `reshape.cpp` 為例，詳細說明如何在 TensorFlow Lite Micro 中添加一個新的操作符。

操作符注冊不僅是模型推理的基礎，更是優化性能、減少內存占用的關鍵環節。掌握這一機制，開發者可以更靈活地定制算子，滿足特定硬件和應用需求。

在 TFLite Micro 中，每個操作符都需要經過以下幾個關鍵步驟：

1. 內核實現：定義操作符的具體計算邏輯

2. 參數解析：從 FlatBuffer 格式中解析操作符參數

3. 操作符注冊：將操作符注冊到解析器中，使其可被模型調用

4. 內存管理：處理張量的內存分配和釋放

操作符實現的核心組件

1. 文件結構說明

添加新操作符需要修改以下幾個關鍵文件，每個文件都有其特定的作用：

micro/kernels/reshape.cpp #

操作符的核心計算邏輯實現

micro/micro_mutable_op_resolver.h#

可變操作符解析器，用于動態注冊操作符

core/api/flatbuffer_conversions.h #

FlatBuffer 參數解析函數的聲明

core/api/flatbuffer_conversions.cpp #

FlatBuffer 參數解析函數的具體實現

micro/all_ops_resolver.cpp #

全局操作符解析器，包含所有支持的操作符

文件作用詳解：

`micro/kernels/` 目錄：

存放所有操作符的具體實現，每個操作符一個文件

`micro_mutable_op_resolver.h`：

提供靈活的操作符注冊接口，允許用戶選擇性地添加操作符

`flatbuffer_conversions.*`：

處理模型文件中的參數解析，將 FlatBuffer 格式轉換為 C++ 結構體

`all_ops_resolver.cpp`：

預定義了所有標準操作符的注冊，適用于需要完整操作符支持的場景

2. 核心實現文件分析

2.1 頭文件引入

文件位置：`micro/kernels/reshape.cpp`

#include
#include"tensorflow/lite/c/builtin_op_data.h"
#include"tensorflow/lite/c/common.h"
#include"tensorflow/lite/kernels/internal/tensor_ctypes.h"
#include"tensorflow/lite/kernels/kernel_util.h"
#include"tensorflow/lite/kernels/op_macros.h"
#include"tensorflow/lite/micro/kernels/kernel_util.h"
#include"tensorflow/lite/micro/memory_helpers.h"
#include"tensorflow/lite/micro/micro_utils.h"

頭文件說明：

`builtin_op_data.h`：包含所有內置操作符的參數結構體定義

`common.h`：TFLite 的基礎數據類型和狀態碼定義

`tensor_ctypes.h`：張量數據類型相關的工具函數

`kernel_util.h`：操作符實現的通用工具函數

`op_macros.h`：操作符實現中常用的宏定義

`micro/kernels/kernel_util.h`：Micro 版本特有的內核工具函數

`memory_helpers.h`：內存管理相關的輔助函數

`micro_utils.h`：Micro 版本的通用工具函數

2.2 命名空間和常量定義

namespacetflite {
namespaceops {
namespacemicro {
namespacereshape {
constexprintkInputTensor =0;
constexprintkOutputTensor =0;

命名空間說明：

`tflite::reshape`：四層命名空間確保了代碼的組織性和避免命名沖突

常量定義：`kInputTensor` 和 `kOutputTensor` 定義了輸入輸出張量的索引，Reshape 操作只有一個輸入和一個輸出

2.3 核心函數實現

ReshapeOutput 函數 - 形狀計算邏輯

TfLiteStatus ReshapeOutput(TfLiteContext* context, TfLiteNode* node) {
 MicroContext* micro_context =GetMicroContext(context);
// 獲取輸入和輸出張量 - 使用臨時分配避免持久內存占用
 TfLiteTensor* input = micro_context->AllocateTempInputTensor(node, kInputTensor);
 TfLiteTensor* output = micro_context->AllocateTempOutputTensor(node, kOutputTensor);
// 計算輸入元素總數 - 用于驗證 reshape 操作的合法性
intnum_input_elements =NumElements(input);
 TfLiteIntArray* output_shape = output->dims;
// 處理特殊情況：-1 維度自動計算
// TensorFlow 允許一個維度設為 -1，表示根據其他維度自動推斷
intnum_output_elements =1;
intstretch_dim = -1;
for(inti =0; i < output_shape->size; ++i) {
 intvalue = output_shape->data[i];
 if(value == -1) {
  TF_LITE_ENSURE_EQ(context, stretch_dim, -1); // 確保只有一個 -1 維度
   stretch_dim = i;
  }else{
   num_output_elements *= value;
  }
 }
// 如果存在 -1 維度，自動計算其大小
if(stretch_dim != -1) {
  TfLiteEvalTensor* output_eval =
    tflite::micro::GetEvalOutput(context, node, kOutputTensor);
 TF_LITE_ENSURE_STATUS(tflite::micro::CreateWritableTensorDimsWithCopy(
    context, output, output_eval));
  output_shape = output->dims; // 更新形狀指針
  output_shape->data[stretch_dim] = num_input_elements / num_output_elements;
  num_output_elements *= output_shape->data[stretch_dim];
 }
// 確保輸入輸出元素數量一致 - Reshape 不改變元素總數
TF_LITE_ENSURE_EQ(context, num_input_elements, num_output_elements);
TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type); // 確保數據類型一致
// 釋放臨時張量 - 避免內存泄漏
 micro_context->DeallocateTempTfLiteTensor(input);
 micro_context->DeallocateTempTfLiteTensor(output);
returnkTfLiteOk;
}

函數作用詳解：

臨時張量分配：使用 `AllocateTempInputTensor` 和 `AllocateTempOutputTensor` 獲取張量信息，這些是臨時分配，不會占用持久內存

形狀驗證：確保 reshape 操作的合法性，輸入輸出元素總數必須相等

自動維度推斷：處理 -1 維度的特殊情況，這是 TensorFlow 的標準特性

內存管理：及時釋放臨時張量，這在內存受限的微控制器環境中非常重要

Prepare 函數-操作符準備階段：

TfLiteStatusPrepare(TfLiteContext* context, TfLiteNode* node){
// 驗證輸入輸出數量 - Reshape 可以有 1 個或 2 個輸入（第二個輸入是可選的形狀參數）
 TF_LITE_ENSURE(context, NumInputs(node) ==1|| NumInputs(node) ==2);
 TF_LITE_ENSURE_EQ(context, NumOutputs(node),1); // 只有一個輸出
// 執行輸出形狀重塑 - 在準備階段確定最終的輸出形狀
 TF_LITE_ENSURE_EQ(context, ReshapeOutput(context, node), kTfLiteOk);
returnkTfLiteOk;
}

Prepare函數說明：

輸入驗證：Reshape 操作支持 1-2 個輸入，第二個輸入是可選的形狀張量

形狀計算：在準備階段就確定輸出形狀，避免在執行階段重復計算

錯誤檢查：使用 `TF_LITE_ENSURE` 宏進行參數驗證，失敗時會返回錯誤狀態

Eval 函數-操作符執行階段：

TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
// 獲取輸入輸出張量 - 使用 EvalTensor 進行實際計算
constTfLiteEvalTensor* input =
   tflite::GetEvalInput(context, node, kInputTensor);
 TfLiteEvalTensor* output =
   tflite::GetEvalOutput(context, node, kOutputTensor);
// 計算輸入數據大小 - 需要拷貝的字節數
 size_t input_bytes;
 TF_LITE_ENSURE_STATUS(TfLiteTypeSizeOf(input->type, &input_bytes));
 input_bytes *= ElementCount(*input->dims);
// 執行數據拷貝（如果不是原地操作）
// 原地操作：輸入輸出使用同一塊內存，無需拷貝
if(input->data.raw != output->data.raw) {
  memcpy(output->data.raw, input->data.raw, input_bytes);
 }
returnkTfLiteOk;
}

Eval函數說明：

EvalTensor 使用：在執行階段使用 `TfLiteEvalTensor`，它包含實際的數據指針

原地操作優化：檢查輸入輸出是否共享內存，避免不必要的數據拷貝

內存拷貝：Reshape 操作本質上只是改變數據的解釋方式，不改變數據內容

2.4 操作符注冊函數

TfLiteRegistration_V1 Register_RESHAPE() {
returntflite::micro::RegisterOp(nullptr, reshape::Prepare, reshape::Eval);
}
} // namespace reshape
} // namespace micro
} // namespace ops
} // namespace tflite

注冊函數說明：

RegisterOp 函數：創建操作符注冊結構體，包含初始化、準備和執行函數指針

nullptr 參數：第一個參數是初始化函數，Reshape 不需要特殊初始化，所以傳入 nullptr

函數指針：傳入 Prepare 和 Eval 函數指針，框架會在適當時機調用這些函數

3. 參數解析實現

3.1 解析函數聲明

文件位置：`core/api/flatbuffer_conversions.h`

TfLiteStatusParseReshape(constOperator* op, ErrorReporter* error_reporter,
             BuiltinDataAllocator* allocator,voidbuiltin_data);

聲明說明：

Operator* op：來自 FlatBuffer 的操作符定義，包含所有參數信息

ErrorReporter：用于報告解析過程中的錯誤

BuiltinDataAllocator：專用的內存分配器，用于分配參數結構體

builtin_data：輸出參數，指向解析后的參數結構體

3.2 解析函數實現

文件位置：`core/api/flatbuffer_conversions.cpp`

TfLiteStatusParseReshape(constOperator* op,ErrorReporter* error_reporter,
            BuiltinDataAllocator* allocator,
            voidbuiltin_data) ；

解析函數詳解：

參數驗證：`CheckParsePointerParams` 確保所有指針參數有效

安全分配器：`SafeBuiltinDataAllocator` 提供異常安全的內存分配

FlatBuffer 解析：從序列化的模型文件中提取 reshape 參數

格式轉換：將 FlatBuffer 格式轉換為 TFLite 內部使用的 C 結構體格式

所有權轉移：使用 `release()` 將參數結構體的所有權轉移給框架

3.3 在解析開關中添加對應的case

文件位置：`core/api/flatbuffer_conversions.cpp`

在 `ParseOpData` 函數的 switch 語句中添加：

caseBuiltinOperator_RESHAPE: {
returnParseReshape(op, error_reporter, allocator, builtin_data);
}

開關語句說明：

這個 switch 語句是 TFLite 參數解析的核心分發機制

根據操作符類型調用相應的解析函數

`BuiltinOperator_RESHAPE` 是在FlatBuffer schema中定義的枚舉值

通過本指南，我們深入了解了 TensorFlow Lite Micro 的操作符注冊機制，包括其設計理念、實現方式以及在嵌入式場景中的重要性。

未來，隨著邊緣計算和微控制器 AI 的快速發展，理解并運用這些底層機制將成為構建高效、可擴展 AI 系統的核心能力。建議讀者在實踐中嘗試自定義算子注冊，并結合實際項目進行優化，以真正發揮 TensorFlow Lite Micro 的潛力。