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

操作符注冊不僅是模型推理的基礎(chǔ)，更是優(yōu)化性能、減少內(nèi)存占用的關(guān)鍵環(huán)節(jié)。掌握這一機制，開發(fā)者可以更靈活地定制算子，滿足特定硬件和應(yīng)用需求。

在 TFLite Micro 中，每個操作符都需要經(jīng)過以下幾個關(guān)鍵步驟：

1. 內(nèi)核實現(xiàn)：定義操作符的具體計算邏輯

2. 參數(shù)解析：從 FlatBuffer 格式中解析操作符參數(shù)

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

4. 內(nèi)存管理：處理張量的內(nèi)存分配和釋放

操作符實現(xiàn)的核心組件

1. 文件結(jié)構(gòu)說明

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

micro/kernels/reshape.cpp #

操作符的核心計算邏輯實現(xiàn)

micro/micro_mutable_op_resolver.h#

可變操作符解析器，用于動態(tài)注冊操作符

core/api/flatbuffer_conversions.h #

FlatBuffer 參數(shù)解析函數(shù)的聲明

core/api/flatbuffer_conversions.cpp #

FlatBuffer 參數(shù)解析函數(shù)的具體實現(xiàn)

micro/all_ops_resolver.cpp #

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

文件作用詳解：

`micro/kernels/` 目錄：

存放所有操作符的具體實現(xiàn)，每個操作符一個文件

`micro_mutable_op_resolver.h`：

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

`flatbuffer_conversions.*`：

處理模型文件中的參數(shù)解析，將 FlatBuffer 格式轉(zhuǎn)換為 C++ 結(jié)構(gòu)體

`all_ops_resolver.cpp`：

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

2. 核心實現(xiàn)文件分析

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`：包含所有內(nèi)置操作符的參數(shù)結(jié)構(gòu)體定義

`common.h`：TFLite 的基礎(chǔ)數(shù)據(jù)類型和狀態(tài)碼定義

`tensor_ctypes.h`：張量數(shù)據(jù)類型相關(guān)的工具函數(shù)

`kernel_util.h`：操作符實現(xiàn)的通用工具函數(shù)

`op_macros.h`：操作符實現(xiàn)中常用的宏定義

`micro/kernels/kernel_util.h`：Micro 版本特有的內(nèi)核工具函數(shù)

`memory_helpers.h`：內(nèi)存管理相關(guān)的輔助函數(shù)

`micro_utils.h`：Micro 版本的通用工具函數(shù)

2.2 命名空間和常量定義

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

命名空間說明：

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

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

2.3 核心函數(shù)實現(xiàn)

ReshapeOutput 函數(shù) - 形狀計算邏輯

TfLiteStatus ReshapeOutput(TfLiteContext* context, TfLiteNode* node) {
 MicroContext* micro_context =GetMicroContext(context);
// 獲取輸入和輸出張量 - 使用臨時分配避免持久內(nèi)存占用
 TfLiteTensor* input = micro_context->AllocateTempInputTensor(node, kInputTensor);
 TfLiteTensor* output = micro_context->AllocateTempOutputTensor(node, kOutputTensor);
// 計算輸入元素總數(shù) - 用于驗證 reshape 操作的合法性
intnum_input_elements =NumElements(input);
 TfLiteIntArray* output_shape = output->dims;
// 處理特殊情況：-1 維度自動計算
// TensorFlow 允許一個維度設(shè)為 -1，表示根據(jù)其他維度自動推斷
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];
 }
// 確保輸入輸出元素數(shù)量一致 - Reshape 不改變元素總數(shù)
TF_LITE_ENSURE_EQ(context, num_input_elements, num_output_elements);
TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type); // 確保數(shù)據(jù)類型一致
// 釋放臨時張量 - 避免內(nèi)存泄漏
 micro_context->DeallocateTempTfLiteTensor(input);
 micro_context->DeallocateTempTfLiteTensor(output);
returnkTfLiteOk;
}

函數(shù)作用詳解：

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

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

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

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

Prepare 函數(shù)-操作符準備階段：

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

Prepare函數(shù)說明：

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

形狀計算：在準備階段就確定輸出形狀，避免在執(zhí)行階段重復(fù)計算

錯誤檢查：使用 `TF_LITE_ENSURE` 宏進行參數(shù)驗證，失敗時會返回錯誤狀態(tài)

Eval 函數(shù)-操作符執(zhí)行階段：

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

Eval函數(shù)說明：

EvalTensor 使用：在執(zhí)行階段使用 `TfLiteEvalTensor`，它包含實際的數(shù)據(jù)指針

原地操作優(yōu)化：檢查輸入輸出是否共享內(nèi)存，避免不必要的數(shù)據(jù)拷貝

內(nèi)存拷貝：Reshape 操作本質(zhì)上只是改變數(shù)據(jù)的解釋方式，不改變數(shù)據(jù)內(nèi)容

2.4 操作符注冊函數(shù)

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

注冊函數(shù)說明：

RegisterOp 函數(shù)：創(chuàng)建操作符注冊結(jié)構(gòu)體，包含初始化、準備和執(zhí)行函數(shù)指針

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

函數(shù)指針：傳入 Prepare 和 Eval 函數(shù)指針，框架會在適當時機調(diào)用這些函數(shù)

3. 參數(shù)解析實現(xiàn)

3.1 解析函數(shù)聲明

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

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

聲明說明：

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

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

BuiltinDataAllocator：專用的內(nèi)存分配器，用于分配參數(shù)結(jié)構(gòu)體

builtin_data：輸出參數(shù)，指向解析后的參數(shù)結(jié)構(gòu)體

3.2 解析函數(shù)實現(xiàn)

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

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

解析函數(shù)詳解：

參數(shù)驗證：`CheckParsePointerParams` 確保所有指針參數(shù)有效

安全分配器：`SafeBuiltinDataAllocator` 提供異常安全的內(nèi)存分配

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

格式轉(zhuǎn)換：將 FlatBuffer 格式轉(zhuǎn)換為 TFLite 內(nèi)部使用的 C 結(jié)構(gòu)體格式

所有權(quán)轉(zhuǎn)移：使用 `release()` 將參數(shù)結(jié)構(gòu)體的所有權(quán)轉(zhuǎn)移給框架

3.3 在解析開關(guān)中添加對應(yīng)的case

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

在 `ParseOpData` 函數(shù)的 switch 語句中添加：

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

開關(guān)語句說明：

這個 switch 語句是 TFLite 參數(shù)解析的核心分發(fā)機制

根據(jù)操作符類型調(diào)用相應(yīng)的解析函數(shù)

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

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

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