與Delegate接管機(jī)制：從FindOp到執(zhí)行計(jì)劃替換全解析)
如果你在移動(dòng)端或者嵌入式設(shè)備上跑過(guò) TensorFlow Lite大概率見(jiàn)過(guò)下面這類(lèi)報(bào)錯(cuò)Didnt find op for builtin opcode BATCH_MATMUL version 3或者遇到更費(fèi)解的情況模型在 PC 上推理一切正常換到某個(gè)硬件板子上就提示 custom op 找不到再或者你興致勃勃接了一個(gè)硬件 delegate結(jié)果日志顯示一個(gè)算子都沒(méi)被接管推理速度紋絲不動(dòng)。這些問(wèn)題背后指向的是同一個(gè)機(jī)制TFLite 的算子注冊(cè)機(jī)制。你想真正定位這類(lèi)問(wèn)題就得順著FindOp這條路一直摸到Delegate。這篇文章我會(huì)從模型里算子的存放方式講起一路拆到OpResolver的查找邏輯再講清楚 delegate 接管執(zhí)行計(jì)劃時(shí)到底發(fā)生了什么最后給一個(gè)可以跑的最小 custom delegate 示例和一套排錯(cuò)思路。內(nèi)容包括算子模型文件結(jié)構(gòu)、TfLiteRegistration內(nèi)核接口、FindOp查找路徑、版本匹配邏輯、BuiltinOpResolver與MutableOpResolver的使用場(chǎng)景、ReplaceNodeSubsetsWithDelegateKernels的執(zhí)行鏈以及手寫(xiě) delegate 的常見(jiàn)坑。適合在部署 TFLite 模型、接入 GPU/NPU 加速、或者要寫(xiě)自定義算子的人看。1. 先明確“算子”在 TFLite 里的三個(gè)身份模型描述、運(yùn)行時(shí)節(jié)點(diǎn)、執(zhí)行內(nèi)核很多人在排查算子問(wèn)題時(shí)會(huì)卡住是因?yàn)闆](méi)分清楚“算子”這個(gè)詞在不同階段指的是不同東西。模型文件里有一個(gè) operator加載進(jìn)解釋器后它變成一個(gè)執(zhí)行節(jié)點(diǎn)真正運(yùn)算時(shí)它又對(duì)應(yīng)一份內(nèi)核代碼。這三個(gè)身份是同一份數(shù)據(jù)在不同環(huán)節(jié)的投影理解它們的對(duì)應(yīng)關(guān)系后面所有問(wèn)題都好辦了。1.1 模型文件里算子是怎么存放的TFLite 模型是 flatbuffer 格式。整個(gè)模型頂層有一張operator_codes表這張表可以理解為“算子字典”table OperatorCode { builtin_code: BuiltinOperator; custom_code: string; version: int; }每個(gè) SubGraph 里有operators數(shù)組數(shù)組里每一個(gè)Operator都通過(guò)opcode_index指向operator_codes里的某一個(gè)條目同時(shí)記錄自己的輸入輸出張量索引table Operator { opcode_index: uint; inputs: [int]; outputs: [int]; }這種設(shè)計(jì)最直觀的意義是省空間一個(gè)模型里哪怕用了 50 次 ADDoperator_codes表里也只存一條 ADD 描述50 個(gè)算子節(jié)點(diǎn)都指向它。更關(guān)鍵的是版本信息只存一份——所有同類(lèi)型算子在轉(zhuǎn)換時(shí)會(huì)被統(tǒng)一寫(xiě)成一個(gè)版本號(hào)。這里有個(gè)容易忽略的點(diǎn)builtin_code是枚舉值custom_code是字符串。內(nèi)置算子走枚舉自定義算子走字符串TFLite 運(yùn)行時(shí)查找這兩類(lèi)算子的方式完全不同后面我會(huì)針對(duì)這一點(diǎn)展開(kāi)。1.2 加載模型后每個(gè)節(jié)點(diǎn)都要“點(diǎn)名”當(dāng)你創(chuàng)建Interpreter時(shí)必須傳入一個(gè)OpResolvertflite::InterpreterBuilder(/* model */, resolver)(interpreter);這個(gè) resolver 就是整本“算子花名冊(cè)”。Interpreter 在初始化階段會(huì)遍歷每個(gè) subgraph 的每個(gè) operator拿著模型文件里的算子描述去 resolver 里“點(diǎn)名”——找對(duì)應(yīng)的內(nèi)核注冊(cè)信息。點(diǎn)名失敗整個(gè)模型加載就會(huì)失敗錯(cuò)誤信息形如Didnt find op for builtin opcode X version Y registration failed一個(gè)容易被忽視的細(xì)節(jié)是點(diǎn)名發(fā)生在Prepare階段之前。也就是說(shuō)即使某個(gè)算子參數(shù)完全合法、輸入輸出形狀也配得上只要 resolver 里沒(méi)有它的注冊(cè)項(xiàng)模型就跑不起來(lái)。注冊(cè)表決定了解釋器“認(rèn)識(shí)”哪些算子而不是“會(huì)算”哪些算子。1.3 TfLiteRegistration四個(gè)函數(shù)指針就是內(nèi)核的全部resolver 里查到的注冊(cè)信息類(lèi)型是TfLiteRegistration。結(jié)構(gòu)主體是四個(gè)函數(shù)指針typedef struct TfLiteRegistration { void* (*init)(TfLiteContext* context, const char* buffer, size_t length); void (*free)(TfLiteContext* context, void* buffer); TfLiteStatus (*prepare)(TfLiteContext* context, TfLiteNode* node); TfLiteStatus (*invoke)(TfLiteContext* context, TfLiteNode* node); int32_t builtin_code; const char* custom_name; int version; } TfLiteRegistration;用生活化的方式理解這四個(gè)函數(shù)init給這個(gè)算子實(shí)例分配私有狀態(tài)相當(dāng)于入職時(shí)領(lǐng)取工位和電腦。free銷(xiāo)毀狀態(tài)相當(dāng)于離職時(shí)歸還設(shè)備。prepare根據(jù)輸入張量形狀推導(dǎo)輸出張量形狀為真正的計(jì)算排好班。invoke執(zhí)行實(shí)際計(jì)算相當(dāng)于正式干活。以 ADD 為例prepare會(huì)讀取輸入張量的 shape給輸出張量也分配同樣的 shapeinvoke才真正逐元素相加。模型里一條builtin_code kTfLiteBuiltinAdd的算子運(yùn)行時(shí)對(duì)應(yīng)到這樣一份TfLiteRegistration四個(gè)函數(shù)指針指向 ADD 內(nèi)核的不同實(shí)現(xiàn)函數(shù)。所以在排查算子問(wèn)題時(shí)我習(xí)慣先問(wèn)一個(gè)問(wèn)題問(wèn)題出在“花名冊(cè)里沒(méi)這個(gè)人”還是“這個(gè)人能力不行 prepare 失敗”還是“干活時(shí)踩坑 invoke 出錯(cuò)”三類(lèi)問(wèn)題的報(bào)錯(cuò)位置和排查手段完全不同。搞清楚這一點(diǎn)比一頭扎進(jìn)源碼里翻找有效得多。2. 順著 FindOp 走一遍內(nèi)置算子的數(shù)組表、自定義算子的哈希表、版本匹配邏輯點(diǎn)名動(dòng)作的核心就是FindOp。它不是一個(gè)普通函數(shù)而是OpResolver基類(lèi)里定義的兩個(gè)虛接口分別應(yīng)對(duì)內(nèi)置算子和自定義算子class OpResolver { public: virtual ~OpResolver() {} virtual const TfLiteRegistration* FindOp(BuiltinOperator op, int version) const 0; virtual const TfLiteRegistration* FindOp(const char* custom_op, int version) const 0; };注意這里有個(gè)容易誤解的點(diǎn)FindOp的返回值是一個(gè)注冊(cè)結(jié)構(gòu)體的指針。解釋器拿這個(gè)指針去調(diào)用對(duì)應(yīng)的函數(shù)而不是自己復(fù)制一份代碼。這也意味著如果 resolver 在運(yùn)行期間生命周期提前結(jié)束指針懸空會(huì)導(dǎo)致崩潰。Android 的 JNI 封裝里如果沒(méi)有把 resolver 和 interpreter 綁定好經(jīng)常會(huì)出現(xiàn)這種“偶發(fā)段錯(cuò)誤”。2.1 內(nèi)置算子的查找路徑builtin_code 當(dāng)數(shù)組下標(biāo)BuiltinOpResolver是使用頻率最高的 resolver 實(shí)現(xiàn)它的內(nèi)部組織方式很簡(jiǎn)單粗暴——一張按BuiltinOperator枚舉值索引的靜態(tài)數(shù)組或者一組按枚舉值組織的注冊(cè)表。查找內(nèi)置算子時(shí)邏輯大致如下const TfLiteRegistration* BuiltinOpResolver::FindOp( BuiltinOperator op, int version) const { // 按枚舉值查表再校驗(yàn)版本 const TfLiteRegistration* registration LookupBuiltin(op); if (!registration) return nullptr; if (registration-version ! version) return nullptr; return registration; }也就是說(shuō)內(nèi)置算子查找的核心是兩個(gè)匹配條件builtin_code枚舉值相等version版本相等。這里分享一個(gè)實(shí)操經(jīng)驗(yàn)不同 TFLite 版本的BuiltinOperator枚舉值不是穩(wěn)定的。舊版運(yùn)行時(shí)拿到新版轉(zhuǎn)換器生成的模型很可能在枚舉值重排后指向了錯(cuò)誤的注冊(cè)項(xiàng)或者直接查不到。所以我從不在生產(chǎn)環(huán)境里做“TFLite 運(yùn)行時(shí)版本比模型轉(zhuǎn)換版本低一點(diǎn)點(diǎn)”這種將就——寧可升級(jí)依賴(lài)也不要賭枚舉值沒(méi)變。2.2 版本匹配為什么經(jīng)常被忽略O(shè)peratorCode里有version字段TfLiteRegistration里也有version字段。查找時(shí)解釋器會(huì)把模型文件里的版本號(hào)傳給FindOpresolver 內(nèi)部再做比對(duì)。同一個(gè)算子有多個(gè)版本通常意味著行為有細(xì)微差異。比如某些算子新版支持了廣播、或者補(bǔ)了精度問(wèn)題、或者換了更優(yōu)的計(jì)算策略。模型轉(zhuǎn)換器會(huì)根據(jù)模型的實(shí)際使用方式選擇一個(gè)版本號(hào)寫(xiě)進(jìn)文件而運(yùn)行時(shí)的注冊(cè)項(xiàng)也有自己的版本號(hào)。兩者對(duì)不上就報(bào)Didnt find op for builtin opcode MUL version 3這里的 “version 3” 指的是模型里期望的算子版本。報(bào)錯(cuò)含義是resolver 里能找到kTfLiteBuiltinMul的注冊(cè)項(xiàng)但找不到version 3的那個(gè)。一個(gè)常被踩的坑是高版本 convert 出來(lái)的模型拿到低版本 TFLite 上運(yùn)行。新版框架可能因?yàn)橹С至诵滤阕诱Z(yǔ)義就把某個(gè)算子的默認(rèn)版本號(hào)抬高了舊運(yùn)行時(shí)沒(méi)注冊(cè)這個(gè)版本直接拒絕加載。排查這類(lèi)問(wèn)題最直接的辦法查一下當(dāng)前 TFLite 版本對(duì)應(yīng)的算子版本映射表或者干脆把tflite依賴(lài)升級(jí)到和模型轉(zhuǎn)換環(huán)境一致的版本。2.3 自定義算子為什么走字符串匹配自定義算子在模型文件里沒(méi)有枚舉值可用只能靠custom_code字符串標(biāo)識(shí)。FindOp(const char* custom_op, int version)的查找路徑本質(zhì)就是一次unordered_map的字符串查找auto it custom_ops_.find(std::string(custom_op)); if (it custom_ops_.end()) return nullptr; if (it-second.version ! version) return nullptr; return it-second;和內(nèi)置算子最大的區(qū)別在于字符串是精確匹配大小寫(xiě)敏感猶豫一點(diǎn)都不行。轉(zhuǎn)換腳本里寫(xiě)的名字是MyCustomOp注冊(cè)時(shí)寫(xiě)的mycustomop結(jié)果就是找不到。很多人問(wèn)為什么自定義算子的報(bào)錯(cuò)信息里沒(méi)有給出版本不匹配的提示而是直接說(shuō) “Didnt find custom op”。因?yàn)閡nordered_map只按字符串找字符串都沒(méi)命中版本號(hào)自然沒(méi)機(jī)會(huì)參與比較。所以排查自定義算子問(wèn)題時(shí)第一件事永遠(yuǎn)是確認(rèn)模型里的字符串和注冊(cè)時(shí)的字符串一字不差。另外注冊(cè)自定義算子用的接口通常是resolver.AddCustom(MyCustomOp, custom_registration, 1);第三個(gè)參數(shù)就是版本號(hào)。如果你后續(xù)改了自定義算子的實(shí)現(xiàn)并提升了版本號(hào)舊模型兼容性會(huì)立刻下降轉(zhuǎn)換新模型時(shí)也要注意保持寫(xiě)進(jìn)模型的版本和注冊(cè)版本一致。3. OpResolver 這套抽象的實(shí)際價(jià)值裁剪、替換和動(dòng)態(tài)注冊(cè)看到這里你可能會(huì)問(wèn)為什么 TFLite 不直接把所有算子都內(nèi)置到解釋器里非要繞一圈通過(guò) resolver 去找答案藏在一個(gè)現(xiàn)實(shí)需求里TFLite 的目標(biāo)環(huán)境太碎了。從手機(jī)到單片機(jī)從幾百兆內(nèi)存到幾百 KB 內(nèi)存的 MCU全量算子對(duì)服務(wù)端框架沒(méi)問(wèn)題對(duì)端側(cè)嵌入式環(huán)境就是災(zāi)難。注冊(cè)機(jī)制的價(jià)值在于把“解釋器核心”和“算子實(shí)現(xiàn)”解耦讓上層按需攜帶、按需替換。3.1 BuiltinOpResolver 和 MutableOpResolver 的差別BuiltinOpResolver就是前面說(shuō)的“全量花名冊(cè)”所有 TFLite 內(nèi)置算子都注冊(cè)在里面。好處是省心壞處是二進(jìn)制體積大——如果你只需要 MINIMAL 推理背上全套算子顯然吃虧。MutableOpResolver是運(yùn)行時(shí)可變的 resolver支持AddBuiltin和AddCustom動(dòng)態(tài)增加注冊(cè)項(xiàng)。兩者對(duì)比如下項(xiàng)目BuiltinOpResolverMutableOpResolver注冊(cè)范圍編譯期間全量?jī)?nèi)置算子運(yùn)行期按需添加自定義算子需要繼承后 override 或配合使用直接 AddCustom二進(jìn)制體積較大只包含實(shí)際注冊(cè)的內(nèi)核適合場(chǎng)景原型驗(yàn)證、通用部署裁剪包體、插件化架構(gòu)實(shí)際項(xiàng)目里我更多是組合使用先用BuiltinOpResolver兜底再額外AddCustom自己寫(xiě)的算子。但如果是做嚴(yán)格裁剪的固件就會(huì)自己繼承OpResolver只暴露模型里真正出現(xiàn)的那幾個(gè)算子。3.2 裁剪二進(jìn)制體積的實(shí)際姿勢(shì)假設(shè)你的模型只有 ADD、CONV_2D、RELU那完全可以寫(xiě)一個(gè)精簡(jiǎn) resolverclass LiteResolver : public tflite::OpResolver { public: LiteResolver() { AddBuiltin(tflite::BuiltinOperator_ADD, tflite::ops::builtin::Register_ADD()); AddBuiltin(tflite::BuiltinOperator_CONV_2D, tflite::ops::builtin::Register_CONV_2D()); AddBuiltin(tflite::BuiltinOperator_RELU, tflite::ops::builtin::Register_RELU()); } const TfLiteRegistration* FindOp(BuiltinOperator op, int version) const override { return GetBuiltinRegistration(op, version); } const TfLiteRegistration* FindOp(const char* custom_op, int version) const override { return GetCustomRegistration(custom_op, version); } };這個(gè)思路再加一層編譯選項(xiàng)配合內(nèi)核源碼只編譯需要的目標(biāo)文件能明顯壓縮體積。關(guān)鍵是你要先知道模型里到底用了哪些算子——?jiǎng)e靠猜直接寫(xiě)個(gè)小腳本遍歷model.operator_codes打印出來(lái)就行。3.3 和 FlexDelegate 的配合算子在 TFLite 和 TensorFlow 之間銜接還有一種情況模型里混了 TensorFlow 算子和 TFLite 算子。TFLite 轉(zhuǎn)換器遇到不支持的標(biāo)準(zhǔn) TF 算子時(shí)如果打開(kāi)了allow_custom_ops或經(jīng)過(guò)一定配置可能會(huì)把它保留成自定義算子名字通常帶Flex前綴。這些 Flex 算子不會(huì)被BuiltinOpResolver找到需要專(zhuān)門(mén)的FlexDelegate來(lái)接管。這個(gè) delegate 本質(zhì)上還是一個(gè)通過(guò)自定義算子名注冊(cè)的機(jī)制——解釋器先通過(guò) custom op 的字符串把它標(biāo)記出來(lái)再由 delegate 在運(yùn)行時(shí)調(diào)用對(duì)應(yīng)的 TensorFlow Lite Flex 內(nèi)核。所以嚴(yán)格來(lái)說(shuō)一個(gè)模型里可以有三種算子來(lái)源純內(nèi)置算子、純自定義算子、由 delegate 支持的算子。理解FindOp只能解決前兩種遇到第三種時(shí)要看 delegate 的接管路徑這正是下一節(jié)的重點(diǎn)。4. Delegate 接管執(zhí)行計(jì)劃的完整邏輯ModifyGraphWithDelegate 到節(jié)點(diǎn)替換Delegate是 TFLite 里被誤解最多的機(jī)制之一。很多人以為 delegate 是“繞過(guò) FindOp 直接走硬件”這個(gè)說(shuō)法不準(zhǔn)確。準(zhǔn)確的理解是delegate 在 FindOp 之后把已經(jīng)解析好的節(jié)點(diǎn)子圖從執(zhí)行計(jì)劃里摘出來(lái)交給另一個(gè)內(nèi)核執(zhí)行。4.1 從 ModifyGraphWithDelegate 開(kāi)始的調(diào)用鏈常規(guī)接入 delegate 的代碼長(zhǎng)這樣TfLiteDelegate* delegate CreateMyDelegate(); interpreter-ModifyGraphWithDelegate(delegate);ModifyGraphWithDelegate內(nèi)部會(huì)按順序做幾件事遍歷當(dāng)前執(zhí)行計(jì)劃里的所有節(jié)點(diǎn)。調(diào)用 delegate 的Prepare回調(diào)。Prepare內(nèi)部決定要接管哪些節(jié)點(diǎn)并調(diào)用核心替換函數(shù)。TFLite 把被接管節(jié)點(diǎn)重構(gòu)成一個(gè)或多個(gè) delegate kernel 節(jié)點(diǎn)。后續(xù)執(zhí)行時(shí)遇到 delegate kernel 節(jié)點(diǎn)就調(diào)用 delegate 內(nèi)核的invoke。這里的“執(zhí)行計(jì)劃”可不是模型文件里的算子順序。TFLite 在內(nèi)部會(huì)做張量生命周期優(yōu)化、內(nèi)存復(fù)用、節(jié)點(diǎn)重排GetExecutionPlan拿到的節(jié)點(diǎn)順序可能和模型里的 operator 順序不一致。寫(xiě)過(guò) delegate 的人多半都踩過(guò)這個(gè)坑你按模型里的 operator 順序去對(duì)接管節(jié)點(diǎn)結(jié)果發(fā)現(xiàn)執(zhí)行計(jì)劃里的節(jié)點(diǎn)編號(hào)完全對(duì)不上。4.2 TfLiteDelegate 和 Prepare 回調(diào)delegate 本身是一個(gè)結(jié)構(gòu)體關(guān)鍵字段和函數(shù)指針如下略去平臺(tái)相關(guān)字段typedef struct TfLiteDelegate { void* data_; TfLiteStatus (*Prepare)(TfLiteContext* context, TfLiteDelegate* delegate); // ... buffer handle 相關(guān)函數(shù)指針 } TfLiteDelegate;Prepare是整個(gè) delegate 的靈魂。TFLite 執(zhí)行ModifyGraphWithDelegate時(shí)會(huì)回調(diào)它而它要做兩件事決定接管哪些節(jié)點(diǎn)、調(diào)用替換函數(shù)把節(jié)點(diǎn)子圖換掉。TfLiteContext提供了兩個(gè)關(guān)鍵接口用于遍歷節(jié)點(diǎn)TF_LITE_ENSURE_STATUS(context-GetExecutionPlan(context, execution_plan)); TF_LITE_ENSURE_STATUS(context-GetNodeAndRegistration( context, node_index, node, registration));拿到node和registration之后registration-builtin_code或registration-custom_name就是判斷是否該接管的依據(jù)。比如想接管 ADD就判斷registration-builtin_code kTfLiteBuiltinAdd。4.3 ReplaceNodeSubsetsWithDelegateKernels 是真正的開(kāi)關(guān)判定完節(jié)點(diǎn)后最核心的一步是調(diào)用context-ReplaceNodeSubsetsWithDelegateKernels( context, delegate_kernel_registration, nodes_to_replace, delegate);nodes_to_replace是一個(gè)整數(shù)數(shù)組元素是執(zhí)行計(jì)劃里的節(jié)點(diǎn)下標(biāo)。這個(gè)函數(shù)做的事情可以理解為T(mén)FLite 拿著這份名單把節(jié)點(diǎn)集合重新組合成一個(gè)或多個(gè)連通的子圖然后每個(gè)子圖變成一個(gè)“delegate kernel”節(jié)點(diǎn)插入執(zhí)行計(jì)劃。被替換之后原算子的TfLiteRegistration不會(huì)被銷(xiāo)毀它的 inputs、outputs、原始注冊(cè)信息仍然保留在模型運(yùn)行時(shí)數(shù)據(jù)結(jié)構(gòu)里。但它的invoke不會(huì)在 CPU 內(nèi)核路徑上被調(diào)用了——執(zhí)行計(jì)劃已經(jīng)指向 delegate kernel 的注冊(cè)信息后續(xù)跑的是你傳入的delegate_kernel_registration.invoke。這里有個(gè)容易誤會(huì)的點(diǎn)delegate kernel 的invoke不是逐算子調(diào)用的而是按子圖調(diào)用的。如果你接管的子圖里有 10 個(gè)算子你的invoke會(huì)被調(diào)用一次內(nèi)部需要負(fù)責(zé)把這 10 個(gè)算子的計(jì)算統(tǒng)一調(diào)度到硬件后端。這也是為什么 delegate 能跨算子做融合優(yōu)化——它看到了整塊子圖可以做算子融合、緩沖區(qū)復(fù)用而不只是把單個(gè)算子搬到別的硬件上執(zhí)行。4.4 真實(shí)項(xiàng)目里 delegate 的常規(guī)用法最常見(jiàn)的三個(gè) delegate正好代表了三種不同的接入方式Delegate覆蓋范圍典型用法NNAPIAndroid 上的 CPU/GPU/DSP/NPUtflite::StatefulNnapiDelegate delegate(options);GPU delegateiOS/Android 上浮點(diǎn)模型整圖加速TfLiteGpuDelegateV2Create(options);XNNPACK浮點(diǎn)算子的 CPU 優(yōu)化通過(guò) interpreter options 自動(dòng)啟用以 NNAPI 為例簡(jiǎn)單接入是這樣#include tensorflow/lite/delegates/nnapi/nnapi_delegate.h tflite::StatefulNnapiDelegate::Options options; tflite::StatefulNnapiDelegate delegate tflite::StatefulNnapiDelegate(options); interpreter-ModifyGraphWithDelegate(delegate);而從 TFLite 2.x 之后的版本開(kāi)始XNNPACK delegate 往往在創(chuàng)建 interpreter 時(shí)通過(guò)experimental_op_resolver_type或默認(rèn)設(shè)置就參與進(jìn)來(lái)了甚至不需要手動(dòng)創(chuàng)建 delegate 對(duì)象。這些成熟 delegate 能加速跑通底層依賴(lài)的就是 4.2 和 4.3 說(shuō)的這套機(jī)制。理解透替換鏈路后你會(huì)明白兩個(gè)關(guān)鍵結(jié)論FindOp 不決定 delegate 能否接管某個(gè)算子。delegate 判斷的依據(jù)是TfLiteRegistration.builtin_code/custom_name即使 CPU 內(nèi)核根本不存在delegate 也能在 Prepare 階段把它接管走前提是你的 delegate 后端真的能執(zhí)行它。找得到的算子不一定走 CPU找不到的算子也不一定會(huì)加載失敗。這和“resolver 里有沒(méi)有注冊(cè)”是兩套獨(dú)立邏輯只是在實(shí)際執(zhí)行計(jì)劃里交織在一起。4.5 delegate Prepare 失敗后的策略如果 delegate 在Prepare階段遇到不支持的節(jié)點(diǎn)組合策略TFLite 的處理方式取決于 delegate 自己。有的 delegate 會(huì)在內(nèi)部做回退把部分節(jié)點(diǎn)留在 CPU 執(zhí)行有的干脆整體失敗讓解釋器進(jìn)入錯(cuò)誤狀態(tài)。實(shí)際項(xiàng)目中我見(jiàn)過(guò)最典型的場(chǎng)景模型里混了 float 和 quantized 算子GPU delegate 只支持其中一部分如果設(shè)置成嚴(yán)格模式strictPrepare 階段直接失敗設(shè)置成寬松模式就能部分接管剩下回落到 CPU。這也是為什么“接入了 delegate 但速度沒(méi)提升”不一定是你代碼寫(xiě)錯(cuò)可能只是你允許了 delegate 部分接管。這個(gè)判斷點(diǎn)很重要在動(dòng)代碼之前先確認(rèn) delegate 的 options 配置。5. 手寫(xiě)一個(gè)最小 custom delegate把 ADD 算子從 CPU 內(nèi)核手里接過(guò)來(lái)理論鋪墊夠了現(xiàn)在做一個(gè)能跑的最小 demo寫(xiě)一個(gè)只接管 ADD 算子的 custom delegate。這個(gè) demo 的執(zhí)行邏輯其實(shí)就是用 C 代碼逐元素相加本質(zhì)上和 CPU 內(nèi)置內(nèi)核做的事一樣價(jià)值在于讓你完整看到“節(jié)點(diǎn)匹配、子圖替換、后端調(diào)度”三段流程長(zhǎng)什么樣。5.1 定義 delegate 和 Prepare 回調(diào)// demo_delegate.h #ifndef DEMO_DELEGATE_H_ #define DEMO_DELEGATE_H_ #include tensorflow/lite/c/c_api.h #include tensorflow/lite/c/common.h namespace demo { bool IsAddNode(const TfLiteNode* node, const TfLiteRegistration* registration) { return registration-builtin_code kTfLiteBuiltinAdd; } TfLiteStatus DemoDelegatePrepare(TfLiteContext* context, TfLiteDelegate* delegate) { TfLiteIntArray* execution_plan nullptr; TF_LITE_ENSURE_STATUS(context-GetExecutionPlan(context, execution_plan)); TfLiteIntArray* nodes_to_replace TfLiteIntArrayCreate(execution_plan-size); int num_selected 0; for (int i 0; i execution_plan-size; i) { int node_index execution_plan-data[i]; TfLiteNode* node nullptr; TfLiteRegistration* registration nullptr; TF_LITE_ENSURE_STATUS(context-GetNodeAndRegistration( context, node_index, node, registration)); if (IsAddNode(node, registration)) { nodes_to_replace-data[num_selected] node_index; } } if (num_selected 0) { TfLiteIntArrayFree(nodes_to_replace); return kTfLiteOk; } TfLiteIntArray* selected_nodes TfLiteIntArrayCreate(num_selected); for (int i 0; i num_selected; i) { selected_nodes-data[i] nodes_to_replace-data[i]; } TfLiteIntArrayFree(nodes_to_replace); TfLiteRegistration delegate_kernel_registration {0}; delegate_kernel_registration.init DemoDelegateKernelInit; delegate_kernel_registration.free DemoDelegateKernelFree; delegate_kernel_registration.prepare DemoDelegateKernelPrepare; delegate_kernel_registration.invoke DemoDelegateKernelInvoke; TF_LITE_ENSURE_STATUS(context-ReplaceNodeSubsetsWithDelegateKernels( context, delegate_kernel_registration, selected_nodes, delegate)); TfLiteIntArrayFree(selected_nodes); return kTfLiteOk; } } // namespace demo #endif // DEMO_DELEGATE_H_注意這里冒出了一個(gè)實(shí)踐細(xì)節(jié)nodes_to_replace一開(kāi)始按execution_plan-size分配但實(shí)際選出來(lái)的節(jié)點(diǎn)數(shù)量可能遠(yuǎn)小于它。真正傳給ReplaceNodeSubsetsWithDelegateKernels的數(shù)組必須精確保留“連續(xù)的前 num_selected 個(gè)元素”所以我復(fù)制了一個(gè)緊湊數(shù)組。直接傳原數(shù)組會(huì)讓 TFLite 誤以為尾部那些 0 值也是有效節(jié)點(diǎn)下標(biāo)輕則接管數(shù)量不對(duì)重則在節(jié)點(diǎn)索引校驗(yàn)時(shí)直接崩掉。這個(gè)坑在成熟 delegate 源碼里一般不會(huì)顯眼地寫(xiě)出來(lái)因?yàn)楣俜綄?shí)現(xiàn)的寫(xiě)法往往更簡(jiǎn)潔但新手照著精簡(jiǎn)代碼抄非常容易踩。5.2 DemoDelegateKernel 的三件套被替換后的 delegate kernel 也逃不開(kāi) init / free / prepare / invoke 四個(gè)函數(shù)。我的 demo 里 init 只用來(lái)創(chuàng)建一塊私有狀態(tài)void* DemoDelegateKernelInit(TfLiteContext* context, const char* buffer, size_t length) { return new int(0); // 實(shí)際上不需要狀態(tài)只是演示 } void DemoDelegateKernelFree(TfLiteContext* context, void* buffer) { delete static_castint*(buffer); } TfLiteStatus DemoDelegateKernelPrepare(TfLiteContext* context, TfLiteNode* node) { return kTfLiteOk; } TfLiteStatus DemoDelegateKernelInvoke(TfLiteContext* context, TfLiteNode* node) { const TfLiteTensor* input context-GetTensor(context, node-inputs-data[0]); const TfLiteTensor* input2 context-GetTensor(context, node-inputs-data[1]); TfLiteTensor* output context-GetTensor(context, node-outputs-data[0]); const float* a static_castconst float*(input-data.data); const float* b static_castconst float*(input2-data.data); float* out static_castfloat*(output-data.data); int num_elements 1; for (int i 0; i output-dims-size; i) { num_elements * output-dims-data[i]; } for (int i 0; i num_elements; i) { out[i] a[i] b[i]; } return kTfLiteOk; }嚴(yán)格來(lái)說(shuō)prepare在這里什么都不做是不對(duì)的——正規(guī)實(shí)現(xiàn)應(yīng)該根據(jù)輸入推導(dǎo)輸出 shape但 ADD 的內(nèi)核行為已經(jīng)保證輸入輸出 shape 一致所以 demo 里偷懶可以跑真實(shí)項(xiàng)目中至少要做 shape 一致性校驗(yàn)。要提醒的是node-inputs-data[0]和node-inputs-data[1]是張量索引要用context-GetTensor(context, index)拿到實(shí)際的TfLiteTensor指針。有些 kernel 實(shí)現(xiàn)里會(huì)用context-GetMutableTensor等變體取決于你是否要寫(xiě)數(shù)據(jù)。不要直接在node上解引用張量結(jié)構(gòu)那只是索引數(shù)組。5.3 接入 Interpreter 并驗(yàn)證delegate 定義好之后接入方式非常直接TfLiteDelegate my_delegate {0}; my_delegate.data_ nullptr; my_delegate.Prepare demo::DemoDelegatePrepare; // 創(chuàng)建一個(gè)帶 ADD 的模型然后 tflite::InterpreterBuilder(model, resolver)(interpreter); interpreter-ModifyGraphWithDelegate(my_delegate); interpreter-Invoke();驗(yàn)證是否接管成功最實(shí)用的手段是看執(zhí)行計(jì)劃。你可以在DemoDelegatePrepare里打印num_selected或者在DemoDelegateKernelInvoke里打日志。如果 Invoke 時(shí)打印了你的日志說(shuō)明這條鏈路是真的通了—— delegate kernel 進(jìn)入了執(zhí)行計(jì)劃并且被執(zhí)行器調(diào)度到了。有一點(diǎn)必須說(shuō)清楚工業(yè)級(jí) delegate 的 invoke 絕不會(huì)像我這個(gè) demo 一樣逐個(gè)元素算。真實(shí)接力場(chǎng)景里delegate 的 Prepare 已經(jīng)在本后端申請(qǐng)好內(nèi)存、建立好設(shè)備句柄invoke 階段直接把這些節(jié)點(diǎn)打包成一次硬件提交比如一次性把整塊 tensor 數(shù)據(jù)拷到 GPU再提交一個(gè) command buffer。這個(gè) demo 的價(jià)值在于鏈路演示直接拿去生產(chǎn)環(huán)境一定會(huì)遇到性能反噬因?yàn)閱嗡阕忧袚Q帶來(lái)的設(shè)備調(diào)度開(kāi)銷(xiāo)遠(yuǎn)超一個(gè) ADD 本身的計(jì)算開(kāi)銷(xiāo)。5.4 這個(gè) demo 里最容易栽的三個(gè)坑沒(méi)設(shè)置delegate.data_或者Prepare函數(shù)指針沒(méi)填對(duì)調(diào)用ModifyGraphWithDelegate時(shí)可能直接段錯(cuò)誤。這些字段是 POD 結(jié)構(gòu)體里的函數(shù)指針漏一個(gè)就是調(diào)用空函數(shù)nullptr排查起來(lái)很隱蔽。匹配節(jié)點(diǎn)時(shí)用了模型 operator 序號(hào)而不是執(zhí)行計(jì)劃節(jié)點(diǎn)序號(hào)。請(qǐng)務(wù)必從context-GetExecutionPlan遍歷不要自己去模型文件里數(shù) operators。注冊(cè)了 delegate 但沒(méi)有一個(gè)節(jié)點(diǎn)被接管時(shí)ReplaceNodeSubsetsWithDelegateKernels傳空數(shù)組。這個(gè) demo 里我做了num_selected 0的保護(hù)真實(shí)項(xiàng)目里也要處理這種情況。否則有的 TFLite 版本里會(huì)觸發(fā)斷言。我自己第一次寫(xiě)的時(shí)候在第二個(gè)坑上耗了一個(gè)晚上。原因是模型文件里 ADD 是第 3 個(gè)算子執(zhí)行計(jì)劃里它排在第 17 位中間插入了若干張量記憶化優(yōu)化帶來(lái)的重排。后來(lái)我老老實(shí)實(shí)打了節(jié)點(diǎn)下標(biāo)映射關(guān)系才發(fā)現(xiàn)自己一直在按錯(cuò)誤編號(hào)匹配。6. 從報(bào)錯(cuò)信息倒推排查注冊(cè)鏈路的斷點(diǎn)在哪個(gè)環(huán)節(jié)最后分享一套排錯(cuò)思路。每次遇到算子相關(guān)的問(wèn)題我習(xí)慣先從報(bào)錯(cuò)信息判斷斷點(diǎn)位置再往下挖。畢竟 TFLite 的報(bào)錯(cuò)文本通常已經(jīng)很明確地告訴了你該看哪里。6.1 “Didnt find op” 類(lèi)報(bào)錯(cuò)的排查清單報(bào)錯(cuò)內(nèi)容排查方向驗(yàn)證手段Didnt find op for builtin opcode X version Ybuiltin_code 枚舉值不匹配或版本不匹配檢查 TFLite 運(yùn)行時(shí)版本確認(rèn) converter 版本與運(yùn)行時(shí)一致Didnt find op for custom op Foo自定義算子字符串不匹配dump 模型 operator_codes逐字節(jié)對(duì)比注冊(cè)名Custom op Foo is not supported模型轉(zhuǎn)換時(shí)未保留該算子轉(zhuǎn)換時(shí)打開(kāi) allow_custom_ops如果確實(shí)需要保留Node number N failed to prepareFindOp 已成功但 prepare 階段出錯(cuò)查看該算子內(nèi)核實(shí)現(xiàn)的 prepare 邏輯遇到 BuiltinOperator 相關(guān)報(bào)錯(cuò)時(shí)我建議先在本地寫(xiě)個(gè)三行腳本打印模型operator_codes的builtin_code和version枚舉值再對(duì)照builtin_op_resolver源碼里注冊(cè)的版本范圍。這一步能排除掉 80% 的“版本不匹配”問(wèn)題。6.2 一個(gè)真實(shí)場(chǎng)景同一份模型在不同設(shè)備上的奇偶問(wèn)題之前有位同學(xué)在項(xiàng)目里遇到的現(xiàn)象是同一個(gè) SSD MobileNet 模型在開(kāi)發(fā)板 A 上跑得好好的換到板子 B 上就報(bào)Didnt find op for builtin opcode VERSION something兩個(gè)板子跑的是同一個(gè)二進(jìn)制版本唯一區(qū)別是板子 B 的系統(tǒng)庫(kù)里意外帶了一個(gè)更舊版本的libtensorflowlite.so于是動(dòng)態(tài)鏈接時(shí)加載到了舊實(shí)現(xiàn)。這類(lèi)問(wèn)題用文本日志排查很容易被忽略因?yàn)闆](méi)有編譯錯(cuò)誤——鏈接時(shí)符號(hào)存在只是行為不一致。解決方式無(wú)非兩點(diǎn)不用動(dòng)態(tài)庫(kù)版本管理依賴(lài)或者把模型轉(zhuǎn)換與運(yùn)行時(shí)版本做成 CI 校驗(yàn)。6.3 delegate 一個(gè)算子都沒(méi)接管的排查順序如果你已經(jīng)接入了 delegate但推理速度沒(méi)有變化需要排查下面幾步確認(rèn) delegate 的 Prepare 真的被調(diào)用了。在 Prepare 函數(shù)頭尾打日志確認(rèn)不是你的代碼里根本沒(méi)創(chuàng)建 delegate 對(duì)象或用錯(cuò)實(shí)例。確認(rèn) target 節(jié)點(diǎn)真的存在。打印執(zhí)行計(jì)劃里的節(jié)點(diǎn)數(shù)、每個(gè)節(jié)點(diǎn)的 builtin_code 列表和你的匹配條件逐一比對(duì)。特別檢查 quantized 算子很多 delegate 只接管 float 算子模型是 quantized 時(shí)自然一個(gè)都不中。確認(rèn) ReplaceNodeSubsetsWithDelegateKernels 的返回狀態(tài)。如果返回kTfLiteError后續(xù)就不會(huì)有 delegate kernel 節(jié)點(diǎn)。確認(rèn) delegate kernel 的 invoke 真的被調(diào)了。在最外層 invoke 打日志如果沒(méi)日志說(shuō)明執(zhí)行計(jì)劃里根本不存在你的 delegate kernel。確認(rèn)沒(méi)有其他 delegate 搶先接管了同一批節(jié)點(diǎn)。多個(gè) delegate 疊加時(shí)先執(zhí)行的 delegate 可能已經(jīng)把 ADD 節(jié)點(diǎn)替換掉了后注冊(cè)的 delegate 自然匹配不到。這套順序看起來(lái)簡(jiǎn)單但執(zhí)行的時(shí)候一定要借助日志而不是靠“我感覺(jué)”。TFLite 編譯時(shí)如果開(kāi)了 verbose log會(huì)在ModifyGraphWithDelegate階段打印節(jié)點(diǎn)替換的詳細(xì)信息沒(méi)開(kāi)日志時(shí)自己在 Prepare 和 kernel 函數(shù)里埋 printf 是最快的。另一個(gè)有價(jià)值的經(jīng)驗(yàn)是delegate 接管率高不等于端到端延遲一定更低。如果被接管節(jié)點(diǎn)夾在大量 CPU 節(jié)點(diǎn)之間tensor 數(shù)據(jù)反復(fù)在 CPU 和硬件后端之間拷貝開(kāi)銷(xiāo)可能抵消掉加速收益。所以做 delegate 優(yōu)化時(shí)我會(huì)把執(zhí)行計(jì)劃畫(huà)出來(lái)重點(diǎn)看能形成多大塊的連續(xù)子圖而不是追求接管的節(jié)點(diǎn)數(shù)量。這也是為什么前面強(qiáng)調(diào)ReplaceNodeSubsetsWithDelegateKernels是按子圖接管——它給了 delegate 做整塊調(diào)度的機(jī)會(huì)而這個(gè)機(jī)會(huì)需要你在 Prepare 里主動(dòng)用好。最后說(shuō)一句我自己的感受。注冊(cè)機(jī)制看懂了之后TFLite 的很多“玄學(xué)”問(wèn)題會(huì)變得特別直白模型文件里的算子描述是一回事解釋器里的內(nèi)核注冊(cè)是另一回事delegate 又是在執(zhí)行計(jì)劃層面的第三回事。三層之間通過(guò)FindOp和ReplaceNodeSubsetsWithDelegateKernels這兩個(gè)關(guān)鍵點(diǎn)串聯(lián)起來(lái)。以后再遇到“明明注冊(cè)了為什么沒(méi)生效”“delegate 接了為什么沒(méi)加速”這類(lèi)問(wèn)題順著報(bào)錯(cuò)信息回到這一層層的鏈路里去定位多半就會(huì)豁然開(kāi)朗。 TFLite算子注冊(cè)與Delegate接管機(jī)制：從FindOp到執(zhí)行計(jì)劃替換全解析
如果你在移動(dòng)端或者嵌入式設(shè)備上跑過(guò) TensorFlow Lite大概率見(jiàn)過(guò)下面這類(lèi)報(bào)錯(cuò)Didnt find op for builtin opcode BATCH_MATMUL version 3或者遇到更費(fèi)解的情況模型在 PC 上推理一切正常換到某個(gè)硬件板子上就提示 custom op 找不到再或者你興致勃勃接了一個(gè)硬件 delegate結(jié)果日志顯示一個(gè)算子都沒(méi)被接管推理速度紋絲不動(dòng)。這些問(wèn)題背后指向的是同一個(gè)機(jī)制TFLite 的算子注冊(cè)機(jī)制。你想真正定位這類(lèi)問(wèn)題就得順著FindOp這條路一直摸到Delegate。這篇文章我會(huì)從模型里算子的存放方式講起一路拆到OpResolver的查找邏輯再講清楚 delegate 接管執(zhí)行計(jì)劃時(shí)到底發(fā)生了什么最后給一個(gè)可以跑的最小 custom delegate 示例和一套排錯(cuò)思路。內(nèi)容包括算子模型文件結(jié)構(gòu)、TfLiteRegistration內(nèi)核接口、FindOp查找路徑、版本匹配邏輯、BuiltinOpResolver與MutableOpResolver的使用場(chǎng)景、ReplaceNodeSubsetsWithDelegateKernels的執(zhí)行鏈以及手寫(xiě) delegate 的常見(jiàn)坑。適合在部署 TFLite 模型、接入 GPU/NPU 加速、或者要寫(xiě)自定義算子的人看。1. 先明確“算子”在 TFLite 里的三個(gè)身份模型描述、運(yùn)行時(shí)節(jié)點(diǎn)、執(zhí)行內(nèi)核很多人在排查算子問(wèn)題時(shí)會(huì)卡住是因?yàn)闆](méi)分清楚“算子”這個(gè)詞在不同階段指的是不同東西。模型文件里有一個(gè) operator加載進(jìn)解釋器后它變成一個(gè)執(zhí)行節(jié)點(diǎn)真正運(yùn)算時(shí)它又對(duì)應(yīng)一份內(nèi)核代碼。這三個(gè)身份是同一份數(shù)據(jù)在不同環(huán)節(jié)的投影理解它們的對(duì)應(yīng)關(guān)系后面所有問(wèn)題都好辦了。1.1 模型文件里算子是怎么存放的TFLite 模型是 flatbuffer 格式。整個(gè)模型頂層有一張operator_codes表這張表可以理解為“算子字典”table OperatorCode { builtin_code: BuiltinOperator; custom_code: string; version: int; }每個(gè) SubGraph 里有operators數(shù)組數(shù)組里每一個(gè)Operator都通過(guò)opcode_index指向operator_codes里的某一個(gè)條目同時(shí)記錄自己的輸入輸出張量索引table Operator { opcode_index: uint; inputs: [int]; outputs: [int]; }這種設(shè)計(jì)最直觀的意義是省空間一個(gè)模型里哪怕用了 50 次 ADDoperator_codes表里也只存一條 ADD 描述50 個(gè)算子節(jié)點(diǎn)都指向它。更關(guān)鍵的是版本信息只存一份——所有同類(lèi)型算子在轉(zhuǎn)換時(shí)會(huì)被統(tǒng)一寫(xiě)成一個(gè)版本號(hào)。這里有個(gè)容易忽略的點(diǎn)builtin_code是枚舉值custom_code是字符串。內(nèi)置算子走枚舉自定義算子走字符串TFLite 運(yùn)行時(shí)查找這兩類(lèi)算子的方式完全不同后面我會(huì)針對(duì)這一點(diǎn)展開(kāi)。1.2 加載模型后每個(gè)節(jié)點(diǎn)都要“點(diǎn)名”當(dāng)你創(chuàng)建Interpreter時(shí)必須傳入一個(gè)OpResolvertflite::InterpreterBuilder(/* model */, resolver)(interpreter);這個(gè) resolver 就是整本“算子花名冊(cè)”。Interpreter 在初始化階段會(huì)遍歷每個(gè) subgraph 的每個(gè) operator拿著模型文件里的算子描述去 resolver 里“點(diǎn)名”——找對(duì)應(yīng)的內(nèi)核注冊(cè)信息。點(diǎn)名失敗整個(gè)模型加載就會(huì)失敗錯(cuò)誤信息形如Didnt find op for builtin opcode X version Y registration failed一個(gè)容易被忽視的細(xì)節(jié)是點(diǎn)名發(fā)生在Prepare階段之前。也就是說(shuō)即使某個(gè)算子參數(shù)完全合法、輸入輸出形狀也配得上只要 resolver 里沒(méi)有它的注冊(cè)項(xiàng)模型就跑不起來(lái)。注冊(cè)表決定了解釋器“認(rèn)識(shí)”哪些算子而不是“會(huì)算”哪些算子。1.3 TfLiteRegistration四個(gè)函數(shù)指針就是內(nèi)核的全部resolver 里查到的注冊(cè)信息類(lèi)型是TfLiteRegistration。結(jié)構(gòu)主體是四個(gè)函數(shù)指針typedef struct TfLiteRegistration { void* (*init)(TfLiteContext* context, const char* buffer, size_t length); void (*free)(TfLiteContext* context, void* buffer); TfLiteStatus (*prepare)(TfLiteContext* context, TfLiteNode* node); TfLiteStatus (*invoke)(TfLiteContext* context, TfLiteNode* node); int32_t builtin_code; const char* custom_name; int version; } TfLiteRegistration;用生活化的方式理解這四個(gè)函數(shù)init給這個(gè)算子實(shí)例分配私有狀態(tài)相當(dāng)于入職時(shí)領(lǐng)取工位和電腦。free銷(xiāo)毀狀態(tài)相當(dāng)于離職時(shí)歸還設(shè)備。prepare根據(jù)輸入張量形狀推導(dǎo)輸出張量形狀為真正的計(jì)算排好班。invoke執(zhí)行實(shí)際計(jì)算相當(dāng)于正式干活。以 ADD 為例prepare會(huì)讀取輸入張量的 shape給輸出張量也分配同樣的 shapeinvoke才真正逐元素相加。模型里一條builtin_code kTfLiteBuiltinAdd的算子運(yùn)行時(shí)對(duì)應(yīng)到這樣一份TfLiteRegistration四個(gè)函數(shù)指針指向 ADD 內(nèi)核的不同實(shí)現(xiàn)函數(shù)。所以在排查算子問(wèn)題時(shí)我習(xí)慣先問(wèn)一個(gè)問(wèn)題問(wèn)題出在“花名冊(cè)里沒(méi)這個(gè)人”還是“這個(gè)人能力不行 prepare 失敗”還是“干活時(shí)踩坑 invoke 出錯(cuò)”三類(lèi)問(wèn)題的報(bào)錯(cuò)位置和排查手段完全不同。搞清楚這一點(diǎn)比一頭扎進(jìn)源碼里翻找有效得多。2. 順著 FindOp 走一遍內(nèi)置算子的數(shù)組表、自定義算子的哈希表、版本匹配邏輯點(diǎn)名動(dòng)作的核心就是FindOp。它不是一個(gè)普通函數(shù)而是OpResolver基類(lèi)里定義的兩個(gè)虛接口分別應(yīng)對(duì)內(nèi)置算子和自定義算子class OpResolver { public: virtual ~OpResolver() {} virtual const TfLiteRegistration* FindOp(BuiltinOperator op, int version) const 0; virtual const TfLiteRegistration* FindOp(const char* custom_op, int version) const 0; };注意這里有個(gè)容易誤解的點(diǎn)FindOp的返回值是一個(gè)注冊(cè)結(jié)構(gòu)體的指針。解釋器拿這個(gè)指針去調(diào)用對(duì)應(yīng)的函數(shù)而不是自己復(fù)制一份代碼。這也意味著如果 resolver 在運(yùn)行期間生命周期提前結(jié)束指針懸空會(huì)導(dǎo)致崩潰。Android 的 JNI 封裝里如果沒(méi)有把 resolver 和 interpreter 綁定好經(jīng)常會(huì)出現(xiàn)這種“偶發(fā)段錯(cuò)誤”。2.1 內(nèi)置算子的查找路徑builtin_code 當(dāng)數(shù)組下標(biāo)BuiltinOpResolver是使用頻率最高的 resolver 實(shí)現(xiàn)它的內(nèi)部組織方式很簡(jiǎn)單粗暴——一張按BuiltinOperator枚舉值索引的靜態(tài)數(shù)組或者一組按枚舉值組織的注冊(cè)表。查找內(nèi)置算子時(shí)邏輯大致如下const TfLiteRegistration* BuiltinOpResolver::FindOp( BuiltinOperator op, int version) const { // 按枚舉值查表再校驗(yàn)版本 const TfLiteRegistration* registration LookupBuiltin(op); if (!registration) return nullptr; if (registration-version ! version) return nullptr; return registration; }也就是說(shuō)內(nèi)置算子查找的核心是兩個(gè)匹配條件builtin_code枚舉值相等version版本相等。這里分享一個(gè)實(shí)操經(jīng)驗(yàn)不同 TFLite 版本的BuiltinOperator枚舉值不是穩(wěn)定的。舊版運(yùn)行時(shí)拿到新版轉(zhuǎn)換器生成的模型很可能在枚舉值重排后指向了錯(cuò)誤的注冊(cè)項(xiàng)或者直接查不到。所以我從不在生產(chǎn)環(huán)境里做“TFLite 運(yùn)行時(shí)版本比模型轉(zhuǎn)換版本低一點(diǎn)點(diǎn)”這種將就——寧可升級(jí)依賴(lài)也不要賭枚舉值沒(méi)變。2.2 版本匹配為什么經(jīng)常被忽略O(shè)peratorCode里有version字段TfLiteRegistration里也有version字段。查找時(shí)解釋器會(huì)把模型文件里的版本號(hào)傳給FindOpresolver 內(nèi)部再做比對(duì)。同一個(gè)算子有多個(gè)版本通常意味著行為有細(xì)微差異。比如某些算子新版支持了廣播、或者補(bǔ)了精度問(wèn)題、或者換了更優(yōu)的計(jì)算策略。模型轉(zhuǎn)換器會(huì)根據(jù)模型的實(shí)際使用方式選擇一個(gè)版本號(hào)寫(xiě)進(jìn)文件而運(yùn)行時(shí)的注冊(cè)項(xiàng)也有自己的版本號(hào)。兩者對(duì)不上就報(bào)Didnt find op for builtin opcode MUL version 3這里的 “version 3” 指的是模型里期望的算子版本。報(bào)錯(cuò)含義是resolver 里能找到kTfLiteBuiltinMul的注冊(cè)項(xiàng)但找不到version 3的那個(gè)。一個(gè)常被踩的坑是高版本 convert 出來(lái)的模型拿到低版本 TFLite 上運(yùn)行。新版框架可能因?yàn)橹С至诵滤阕诱Z(yǔ)義就把某個(gè)算子的默認(rèn)版本號(hào)抬高了舊運(yùn)行時(shí)沒(méi)注冊(cè)這個(gè)版本直接拒絕加載。排查這類(lèi)問(wèn)題最直接的辦法查一下當(dāng)前 TFLite 版本對(duì)應(yīng)的算子版本映射表或者干脆把tflite依賴(lài)升級(jí)到和模型轉(zhuǎn)換環(huán)境一致的版本。2.3 自定義算子為什么走字符串匹配自定義算子在模型文件里沒(méi)有枚舉值可用只能靠custom_code字符串標(biāo)識(shí)。FindOp(const char* custom_op, int version)的查找路徑本質(zhì)就是一次unordered_map的字符串查找auto it custom_ops_.find(std::string(custom_op)); if (it custom_ops_.end()) return nullptr; if (it-second.version ! version) return nullptr; return it-second;和內(nèi)置算子最大的區(qū)別在于字符串是精確匹配大小寫(xiě)敏感猶豫一點(diǎn)都不行。轉(zhuǎn)換腳本里寫(xiě)的名字是MyCustomOp注冊(cè)時(shí)寫(xiě)的mycustomop結(jié)果就是找不到。很多人問(wèn)為什么自定義算子的報(bào)錯(cuò)信息里沒(méi)有給出版本不匹配的提示而是直接說(shuō) “Didnt find custom op”。因?yàn)閡nordered_map只按字符串找字符串都沒(méi)命中版本號(hào)自然沒(méi)機(jī)會(huì)參與比較。所以排查自定義算子問(wèn)題時(shí)第一件事永遠(yuǎn)是確認(rèn)模型里的字符串和注冊(cè)時(shí)的字符串一字不差。另外注冊(cè)自定義算子用的接口通常是resolver.AddCustom(MyCustomOp, custom_registration, 1);第三個(gè)參數(shù)就是版本號(hào)。如果你后續(xù)改了自定義算子的實(shí)現(xiàn)并提升了版本號(hào)舊模型兼容性會(huì)立刻下降轉(zhuǎn)換新模型時(shí)也要注意保持寫(xiě)進(jìn)模型的版本和注冊(cè)版本一致。3. OpResolver 這套抽象的實(shí)際價(jià)值裁剪、替換和動(dòng)態(tài)注冊(cè)看到這里你可能會(huì)問(wèn)為什么 TFLite 不直接把所有算子都內(nèi)置到解釋器里非要繞一圈通過(guò) resolver 去找答案藏在一個(gè)現(xiàn)實(shí)需求里TFLite 的目標(biāo)環(huán)境太碎了。從手機(jī)到單片機(jī)從幾百兆內(nèi)存到幾百 KB 內(nèi)存的 MCU全量算子對(duì)服務(wù)端框架沒(méi)問(wèn)題對(duì)端側(cè)嵌入式環(huán)境就是災(zāi)難。注冊(cè)機(jī)制的價(jià)值在于把“解釋器核心”和“算子實(shí)現(xiàn)”解耦讓上層按需攜帶、按需替換。3.1 BuiltinOpResolver 和 MutableOpResolver 的差別BuiltinOpResolver就是前面說(shuō)的“全量花名冊(cè)”所有 TFLite 內(nèi)置算子都注冊(cè)在里面。好處是省心壞處是二進(jìn)制體積大——如果你只需要 MINIMAL 推理背上全套算子顯然吃虧。MutableOpResolver是運(yùn)行時(shí)可變的 resolver支持AddBuiltin和AddCustom動(dòng)態(tài)增加注冊(cè)項(xiàng)。兩者對(duì)比如下項(xiàng)目BuiltinOpResolverMutableOpResolver注冊(cè)范圍編譯期間全量?jī)?nèi)置算子運(yùn)行期按需添加自定義算子需要繼承后 override 或配合使用直接 AddCustom二進(jìn)制體積較大只包含實(shí)際注冊(cè)的內(nèi)核適合場(chǎng)景原型驗(yàn)證、通用部署裁剪包體、插件化架構(gòu)實(shí)際項(xiàng)目里我更多是組合使用先用BuiltinOpResolver兜底再額外AddCustom自己寫(xiě)的算子。但如果是做嚴(yán)格裁剪的固件就會(huì)自己繼承OpResolver只暴露模型里真正出現(xiàn)的那幾個(gè)算子。3.2 裁剪二進(jìn)制體積的實(shí)際姿勢(shì)假設(shè)你的模型只有 ADD、CONV_2D、RELU那完全可以寫(xiě)一個(gè)精簡(jiǎn) resolverclass LiteResolver : public tflite::OpResolver { public: LiteResolver() { AddBuiltin(tflite::BuiltinOperator_ADD, tflite::ops::builtin::Register_ADD()); AddBuiltin(tflite::BuiltinOperator_CONV_2D, tflite::ops::builtin::Register_CONV_2D()); AddBuiltin(tflite::BuiltinOperator_RELU, tflite::ops::builtin::Register_RELU()); } const TfLiteRegistration* FindOp(BuiltinOperator op, int version) const override { return GetBuiltinRegistration(op, version); } const TfLiteRegistration* FindOp(const char* custom_op, int version) const override { return GetCustomRegistration(custom_op, version); } };這個(gè)思路再加一層編譯選項(xiàng)配合內(nèi)核源碼只編譯需要的目標(biāo)文件能明顯壓縮體積。關(guān)鍵是你要先知道模型里到底用了哪些算子——?jiǎng)e靠猜直接寫(xiě)個(gè)小腳本遍歷model.operator_codes打印出來(lái)就行。3.3 和 FlexDelegate 的配合算子在 TFLite 和 TensorFlow 之間銜接還有一種情況模型里混了 TensorFlow 算子和 TFLite 算子。TFLite 轉(zhuǎn)換器遇到不支持的標(biāo)準(zhǔn) TF 算子時(shí)如果打開(kāi)了allow_custom_ops或經(jīng)過(guò)一定配置可能會(huì)把它保留成自定義算子名字通常帶Flex前綴。這些 Flex 算子不會(huì)被BuiltinOpResolver找到需要專(zhuān)門(mén)的FlexDelegate來(lái)接管。這個(gè) delegate 本質(zhì)上還是一個(gè)通過(guò)自定義算子名注冊(cè)的機(jī)制——解釋器先通過(guò) custom op 的字符串把它標(biāo)記出來(lái)再由 delegate 在運(yùn)行時(shí)調(diào)用對(duì)應(yīng)的 TensorFlow Lite Flex 內(nèi)核。所以嚴(yán)格來(lái)說(shuō)一個(gè)模型里可以有三種算子來(lái)源純內(nèi)置算子、純自定義算子、由 delegate 支持的算子。理解FindOp只能解決前兩種遇到第三種時(shí)要看 delegate 的接管路徑這正是下一節(jié)的重點(diǎn)。4. Delegate 接管執(zhí)行計(jì)劃的完整邏輯ModifyGraphWithDelegate 到節(jié)點(diǎn)替換Delegate是 TFLite 里被誤解最多的機(jī)制之一。很多人以為 delegate 是“繞過(guò) FindOp 直接走硬件”這個(gè)說(shuō)法不準(zhǔn)確。準(zhǔn)確的理解是delegate 在 FindOp 之后把已經(jīng)解析好的節(jié)點(diǎn)子圖從執(zhí)行計(jì)劃里摘出來(lái)交給另一個(gè)內(nèi)核執(zhí)行。4.1 從 ModifyGraphWithDelegate 開(kāi)始的調(diào)用鏈常規(guī)接入 delegate 的代碼長(zhǎng)這樣TfLiteDelegate* delegate CreateMyDelegate(); interpreter-ModifyGraphWithDelegate(delegate);ModifyGraphWithDelegate內(nèi)部會(huì)按順序做幾件事遍歷當(dāng)前執(zhí)行計(jì)劃里的所有節(jié)點(diǎn)。調(diào)用 delegate 的Prepare回調(diào)。Prepare內(nèi)部決定要接管哪些節(jié)點(diǎn)并調(diào)用核心替換函數(shù)。TFLite 把被接管節(jié)點(diǎn)重構(gòu)成一個(gè)或多個(gè) delegate kernel 節(jié)點(diǎn)。后續(xù)執(zhí)行時(shí)遇到 delegate kernel 節(jié)點(diǎn)就調(diào)用 delegate 內(nèi)核的invoke。這里的“執(zhí)行計(jì)劃”可不是模型文件里的算子順序。TFLite 在內(nèi)部會(huì)做張量生命周期優(yōu)化、內(nèi)存復(fù)用、節(jié)點(diǎn)重排GetExecutionPlan拿到的節(jié)點(diǎn)順序可能和模型里的 operator 順序不一致。寫(xiě)過(guò) delegate 的人多半都踩過(guò)這個(gè)坑你按模型里的 operator 順序去對(duì)接管節(jié)點(diǎn)結(jié)果發(fā)現(xiàn)執(zhí)行計(jì)劃里的節(jié)點(diǎn)編號(hào)完全對(duì)不上。4.2 TfLiteDelegate 和 Prepare 回調(diào)delegate 本身是一個(gè)結(jié)構(gòu)體關(guān)鍵字段和函數(shù)指針如下略去平臺(tái)相關(guān)字段typedef struct TfLiteDelegate { void* data_; TfLiteStatus (*Prepare)(TfLiteContext* context, TfLiteDelegate* delegate); // ... buffer handle 相關(guān)函數(shù)指針 } TfLiteDelegate;Prepare是整個(gè) delegate 的靈魂。TFLite 執(zhí)行ModifyGraphWithDelegate時(shí)會(huì)回調(diào)它而它要做兩件事決定接管哪些節(jié)點(diǎn)、調(diào)用替換函數(shù)把節(jié)點(diǎn)子圖換掉。TfLiteContext提供了兩個(gè)關(guān)鍵接口用于遍歷節(jié)點(diǎn)TF_LITE_ENSURE_STATUS(context-GetExecutionPlan(context, execution_plan)); TF_LITE_ENSURE_STATUS(context-GetNodeAndRegistration( context, node_index, node, registration));拿到node和registration之后registration-builtin_code或registration-custom_name就是判斷是否該接管的依據(jù)。比如想接管 ADD就判斷registration-builtin_code kTfLiteBuiltinAdd。4.3 ReplaceNodeSubsetsWithDelegateKernels 是真正的開(kāi)關(guān)判定完節(jié)點(diǎn)后最核心的一步是調(diào)用context-ReplaceNodeSubsetsWithDelegateKernels( context, delegate_kernel_registration, nodes_to_replace, delegate);nodes_to_replace是一個(gè)整數(shù)數(shù)組元素是執(zhí)行計(jì)劃里的節(jié)點(diǎn)下標(biāo)。這個(gè)函數(shù)做的事情可以理解為T(mén)FLite 拿著這份名單把節(jié)點(diǎn)集合重新組合成一個(gè)或多個(gè)連通的子圖然后每個(gè)子圖變成一個(gè)“delegate kernel”節(jié)點(diǎn)插入執(zhí)行計(jì)劃。被替換之后原算子的TfLiteRegistration不會(huì)被銷(xiāo)毀它的 inputs、outputs、原始注冊(cè)信息仍然保留在模型運(yùn)行時(shí)數(shù)據(jù)結(jié)構(gòu)里。但它的invoke不會(huì)在 CPU 內(nèi)核路徑上被調(diào)用了——執(zhí)行計(jì)劃已經(jīng)指向 delegate kernel 的注冊(cè)信息后續(xù)跑的是你傳入的delegate_kernel_registration.invoke。這里有個(gè)容易誤會(huì)的點(diǎn)delegate kernel 的invoke不是逐算子調(diào)用的而是按子圖調(diào)用的。如果你接管的子圖里有 10 個(gè)算子你的invoke會(huì)被調(diào)用一次內(nèi)部需要負(fù)責(zé)把這 10 個(gè)算子的計(jì)算統(tǒng)一調(diào)度到硬件后端。這也是為什么 delegate 能跨算子做融合優(yōu)化——它看到了整塊子圖可以做算子融合、緩沖區(qū)復(fù)用而不只是把單個(gè)算子搬到別的硬件上執(zhí)行。4.4 真實(shí)項(xiàng)目里 delegate 的常規(guī)用法最常見(jiàn)的三個(gè) delegate正好代表了三種不同的接入方式Delegate覆蓋范圍典型用法NNAPIAndroid 上的 CPU/GPU/DSP/NPUtflite::StatefulNnapiDelegate delegate(options);GPU delegateiOS/Android 上浮點(diǎn)模型整圖加速TfLiteGpuDelegateV2Create(options);XNNPACK浮點(diǎn)算子的 CPU 優(yōu)化通過(guò) interpreter options 自動(dòng)啟用以 NNAPI 為例簡(jiǎn)單接入是這樣#include tensorflow/lite/delegates/nnapi/nnapi_delegate.h tflite::StatefulNnapiDelegate::Options options; tflite::StatefulNnapiDelegate delegate tflite::StatefulNnapiDelegate(options); interpreter-ModifyGraphWithDelegate(delegate);而從 TFLite 2.x 之后的版本開(kāi)始XNNPACK delegate 往往在創(chuàng)建 interpreter 時(shí)通過(guò)experimental_op_resolver_type或默認(rèn)設(shè)置就參與進(jìn)來(lái)了甚至不需要手動(dòng)創(chuàng)建 delegate 對(duì)象。這些成熟 delegate 能加速跑通底層依賴(lài)的就是 4.2 和 4.3 說(shuō)的這套機(jī)制。理解透替換鏈路后你會(huì)明白兩個(gè)關(guān)鍵結(jié)論FindOp 不決定 delegate 能否接管某個(gè)算子。delegate 判斷的依據(jù)是TfLiteRegistration.builtin_code/custom_name即使 CPU 內(nèi)核根本不存在delegate 也能在 Prepare 階段把它接管走前提是你的 delegate 后端真的能執(zhí)行它。找得到的算子不一定走 CPU找不到的算子也不一定會(huì)加載失敗。這和“resolver 里有沒(méi)有注冊(cè)”是兩套獨(dú)立邏輯只是在實(shí)際執(zhí)行計(jì)劃里交織在一起。4.5 delegate Prepare 失敗后的策略如果 delegate 在Prepare階段遇到不支持的節(jié)點(diǎn)組合策略TFLite 的處理方式取決于 delegate 自己。有的 delegate 會(huì)在內(nèi)部做回退把部分節(jié)點(diǎn)留在 CPU 執(zhí)行有的干脆整體失敗讓解釋器進(jìn)入錯(cuò)誤狀態(tài)。實(shí)際項(xiàng)目中我見(jiàn)過(guò)最典型的場(chǎng)景模型里混了 float 和 quantized 算子GPU delegate 只支持其中一部分如果設(shè)置成嚴(yán)格模式strictPrepare 階段直接失敗設(shè)置成寬松模式就能部分接管剩下回落到 CPU。這也是為什么“接入了 delegate 但速度沒(méi)提升”不一定是你代碼寫(xiě)錯(cuò)可能只是你允許了 delegate 部分接管。這個(gè)判斷點(diǎn)很重要在動(dòng)代碼之前先確認(rèn) delegate 的 options 配置。5. 手寫(xiě)一個(gè)最小 custom delegate把 ADD 算子從 CPU 內(nèi)核手里接過(guò)來(lái)理論鋪墊夠了現(xiàn)在做一個(gè)能跑的最小 demo寫(xiě)一個(gè)只接管 ADD 算子的 custom delegate。這個(gè) demo 的執(zhí)行邏輯其實(shí)就是用 C 代碼逐元素相加本質(zhì)上和 CPU 內(nèi)置內(nèi)核做的事一樣價(jià)值在于讓你完整看到“節(jié)點(diǎn)匹配、子圖替換、后端調(diào)度”三段流程長(zhǎng)什么樣。5.1 定義 delegate 和 Prepare 回調(diào)// demo_delegate.h #ifndef DEMO_DELEGATE_H_ #define DEMO_DELEGATE_H_ #include tensorflow/lite/c/c_api.h #include tensorflow/lite/c/common.h namespace demo { bool IsAddNode(const TfLiteNode* node, const TfLiteRegistration* registration) { return registration-builtin_code kTfLiteBuiltinAdd; } TfLiteStatus DemoDelegatePrepare(TfLiteContext* context, TfLiteDelegate* delegate) { TfLiteIntArray* execution_plan nullptr; TF_LITE_ENSURE_STATUS(context-GetExecutionPlan(context, execution_plan)); TfLiteIntArray* nodes_to_replace TfLiteIntArrayCreate(execution_plan-size); int num_selected 0; for (int i 0; i execution_plan-size; i) { int node_index execution_plan-data[i]; TfLiteNode* node nullptr; TfLiteRegistration* registration nullptr; TF_LITE_ENSURE_STATUS(context-GetNodeAndRegistration( context, node_index, node, registration)); if (IsAddNode(node, registration)) { nodes_to_replace-data[num_selected] node_index; } } if (num_selected 0) { TfLiteIntArrayFree(nodes_to_replace); return kTfLiteOk; } TfLiteIntArray* selected_nodes TfLiteIntArrayCreate(num_selected); for (int i 0; i num_selected; i) { selected_nodes-data[i] nodes_to_replace-data[i]; } TfLiteIntArrayFree(nodes_to_replace); TfLiteRegistration delegate_kernel_registration {0}; delegate_kernel_registration.init DemoDelegateKernelInit; delegate_kernel_registration.free DemoDelegateKernelFree; delegate_kernel_registration.prepare DemoDelegateKernelPrepare; delegate_kernel_registration.invoke DemoDelegateKernelInvoke; TF_LITE_ENSURE_STATUS(context-ReplaceNodeSubsetsWithDelegateKernels( context, delegate_kernel_registration, selected_nodes, delegate)); TfLiteIntArrayFree(selected_nodes); return kTfLiteOk; } } // namespace demo #endif // DEMO_DELEGATE_H_注意這里冒出了一個(gè)實(shí)踐細(xì)節(jié)nodes_to_replace一開(kāi)始按execution_plan-size分配但實(shí)際選出來(lái)的節(jié)點(diǎn)數(shù)量可能遠(yuǎn)小于它。真正傳給ReplaceNodeSubsetsWithDelegateKernels的數(shù)組必須精確保留“連續(xù)的前 num_selected 個(gè)元素”所以我復(fù)制了一個(gè)緊湊數(shù)組。直接傳原數(shù)組會(huì)讓 TFLite 誤以為尾部那些 0 值也是有效節(jié)點(diǎn)下標(biāo)輕則接管數(shù)量不對(duì)重則在節(jié)點(diǎn)索引校驗(yàn)時(shí)直接崩掉。這個(gè)坑在成熟 delegate 源碼里一般不會(huì)顯眼地寫(xiě)出來(lái)因?yàn)楣俜綄?shí)現(xiàn)的寫(xiě)法往往更簡(jiǎn)潔但新手照著精簡(jiǎn)代碼抄非常容易踩。5.2 DemoDelegateKernel 的三件套被替換后的 delegate kernel 也逃不開(kāi) init / free / prepare / invoke 四個(gè)函數(shù)。我的 demo 里 init 只用來(lái)創(chuàng)建一塊私有狀態(tài)void* DemoDelegateKernelInit(TfLiteContext* context, const char* buffer, size_t length) { return new int(0); // 實(shí)際上不需要狀態(tài)只是演示 } void DemoDelegateKernelFree(TfLiteContext* context, void* buffer) { delete static_castint*(buffer); } TfLiteStatus DemoDelegateKernelPrepare(TfLiteContext* context, TfLiteNode* node) { return kTfLiteOk; } TfLiteStatus DemoDelegateKernelInvoke(TfLiteContext* context, TfLiteNode* node) { const TfLiteTensor* input context-GetTensor(context, node-inputs-data[0]); const TfLiteTensor* input2 context-GetTensor(context, node-inputs-data[1]); TfLiteTensor* output context-GetTensor(context, node-outputs-data[0]); const float* a static_castconst float*(input-data.data); const float* b static_castconst float*(input2-data.data); float* out static_castfloat*(output-data.data); int num_elements 1; for (int i 0; i output-dims-size; i) { num_elements * output-dims-data[i]; } for (int i 0; i num_elements; i) { out[i] a[i] b[i]; } return kTfLiteOk; }嚴(yán)格來(lái)說(shuō)prepare在這里什么都不做是不對(duì)的——正規(guī)實(shí)現(xiàn)應(yīng)該根據(jù)輸入推導(dǎo)輸出 shape但 ADD 的內(nèi)核行為已經(jīng)保證輸入輸出 shape 一致所以 demo 里偷懶可以跑真實(shí)項(xiàng)目中至少要做 shape 一致性校驗(yàn)。要提醒的是node-inputs-data[0]和node-inputs-data[1]是張量索引要用context-GetTensor(context, index)拿到實(shí)際的TfLiteTensor指針。有些 kernel 實(shí)現(xiàn)里會(huì)用context-GetMutableTensor等變體取決于你是否要寫(xiě)數(shù)據(jù)。不要直接在node上解引用張量結(jié)構(gòu)那只是索引數(shù)組。5.3 接入 Interpreter 并驗(yàn)證delegate 定義好之后接入方式非常直接TfLiteDelegate my_delegate {0}; my_delegate.data_ nullptr; my_delegate.Prepare demo::DemoDelegatePrepare; // 創(chuàng)建一個(gè)帶 ADD 的模型然后 tflite::InterpreterBuilder(model, resolver)(interpreter); interpreter-ModifyGraphWithDelegate(my_delegate); interpreter-Invoke();驗(yàn)證是否接管成功最實(shí)用的手段是看執(zhí)行計(jì)劃。你可以在DemoDelegatePrepare里打印num_selected或者在DemoDelegateKernelInvoke里打日志。如果 Invoke 時(shí)打印了你的日志說(shuō)明這條鏈路是真的通了—— delegate kernel 進(jìn)入了執(zhí)行計(jì)劃并且被執(zhí)行器調(diào)度到了。有一點(diǎn)必須說(shuō)清楚工業(yè)級(jí) delegate 的 invoke 絕不會(huì)像我這個(gè) demo 一樣逐個(gè)元素算。真實(shí)接力場(chǎng)景里delegate 的 Prepare 已經(jīng)在本后端申請(qǐng)好內(nèi)存、建立好設(shè)備句柄invoke 階段直接把這些節(jié)點(diǎn)打包成一次硬件提交比如一次性把整塊 tensor 數(shù)據(jù)拷到 GPU再提交一個(gè) command buffer。這個(gè) demo 的價(jià)值在于鏈路演示直接拿去生產(chǎn)環(huán)境一定會(huì)遇到性能反噬因?yàn)閱嗡阕忧袚Q帶來(lái)的設(shè)備調(diào)度開(kāi)銷(xiāo)遠(yuǎn)超一個(gè) ADD 本身的計(jì)算開(kāi)銷(xiāo)。5.4 這個(gè) demo 里最容易栽的三個(gè)坑沒(méi)設(shè)置delegate.data_或者Prepare函數(shù)指針沒(méi)填對(duì)調(diào)用ModifyGraphWithDelegate時(shí)可能直接段錯(cuò)誤。這些字段是 POD 結(jié)構(gòu)體里的函數(shù)指針漏一個(gè)就是調(diào)用空函數(shù)nullptr排查起來(lái)很隱蔽。匹配節(jié)點(diǎn)時(shí)用了模型 operator 序號(hào)而不是執(zhí)行計(jì)劃節(jié)點(diǎn)序號(hào)。請(qǐng)務(wù)必從context-GetExecutionPlan遍歷不要自己去模型文件里數(shù) operators。注冊(cè)了 delegate 但沒(méi)有一個(gè)節(jié)點(diǎn)被接管時(shí)ReplaceNodeSubsetsWithDelegateKernels傳空數(shù)組。這個(gè) demo 里我做了num_selected 0的保護(hù)真實(shí)項(xiàng)目里也要處理這種情況。否則有的 TFLite 版本里會(huì)觸發(fā)斷言。我自己第一次寫(xiě)的時(shí)候在第二個(gè)坑上耗了一個(gè)晚上。原因是模型文件里 ADD 是第 3 個(gè)算子執(zhí)行計(jì)劃里它排在第 17 位中間插入了若干張量記憶化優(yōu)化帶來(lái)的重排。后來(lái)我老老實(shí)實(shí)打了節(jié)點(diǎn)下標(biāo)映射關(guān)系才發(fā)現(xiàn)自己一直在按錯(cuò)誤編號(hào)匹配。6. 從報(bào)錯(cuò)信息倒推排查注冊(cè)鏈路的斷點(diǎn)在哪個(gè)環(huán)節(jié)最后分享一套排錯(cuò)思路。每次遇到算子相關(guān)的問(wèn)題我習(xí)慣先從報(bào)錯(cuò)信息判斷斷點(diǎn)位置再往下挖。畢竟 TFLite 的報(bào)錯(cuò)文本通常已經(jīng)很明確地告訴了你該看哪里。6.1 “Didnt find op” 類(lèi)報(bào)錯(cuò)的排查清單報(bào)錯(cuò)內(nèi)容排查方向驗(yàn)證手段Didnt find op for builtin opcode X version Ybuiltin_code 枚舉值不匹配或版本不匹配檢查 TFLite 運(yùn)行時(shí)版本確認(rèn) converter 版本與運(yùn)行時(shí)一致Didnt find op for custom op Foo自定義算子字符串不匹配dump 模型 operator_codes逐字節(jié)對(duì)比注冊(cè)名Custom op Foo is not supported模型轉(zhuǎn)換時(shí)未保留該算子轉(zhuǎn)換時(shí)打開(kāi) allow_custom_ops如果確實(shí)需要保留Node number N failed to prepareFindOp 已成功但 prepare 階段出錯(cuò)查看該算子內(nèi)核實(shí)現(xiàn)的 prepare 邏輯遇到 BuiltinOperator 相關(guān)報(bào)錯(cuò)時(shí)我建議先在本地寫(xiě)個(gè)三行腳本打印模型operator_codes的builtin_code和version枚舉值再對(duì)照builtin_op_resolver源碼里注冊(cè)的版本范圍。這一步能排除掉 80% 的“版本不匹配”問(wèn)題。6.2 一個(gè)真實(shí)場(chǎng)景同一份模型在不同設(shè)備上的奇偶問(wèn)題之前有位同學(xué)在項(xiàng)目里遇到的現(xiàn)象是同一個(gè) SSD MobileNet 模型在開(kāi)發(fā)板 A 上跑得好好的換到板子 B 上就報(bào)Didnt find op for builtin opcode VERSION something兩個(gè)板子跑的是同一個(gè)二進(jìn)制版本唯一區(qū)別是板子 B 的系統(tǒng)庫(kù)里意外帶了一個(gè)更舊版本的libtensorflowlite.so于是動(dòng)態(tài)鏈接時(shí)加載到了舊實(shí)現(xiàn)。這類(lèi)問(wèn)題用文本日志排查很容易被忽略因?yàn)闆](méi)有編譯錯(cuò)誤——鏈接時(shí)符號(hào)存在只是行為不一致。解決方式無(wú)非兩點(diǎn)不用動(dòng)態(tài)庫(kù)版本管理依賴(lài)或者把模型轉(zhuǎn)換與運(yùn)行時(shí)版本做成 CI 校驗(yàn)。6.3 delegate 一個(gè)算子都沒(méi)接管的排查順序如果你已經(jīng)接入了 delegate但推理速度沒(méi)有變化需要排查下面幾步確認(rèn) delegate 的 Prepare 真的被調(diào)用了。在 Prepare 函數(shù)頭尾打日志確認(rèn)不是你的代碼里根本沒(méi)創(chuàng)建 delegate 對(duì)象或用錯(cuò)實(shí)例。確認(rèn) target 節(jié)點(diǎn)真的存在。打印執(zhí)行計(jì)劃里的節(jié)點(diǎn)數(shù)、每個(gè)節(jié)點(diǎn)的 builtin_code 列表和你的匹配條件逐一比對(duì)。特別檢查 quantized 算子很多 delegate 只接管 float 算子模型是 quantized 時(shí)自然一個(gè)都不中。確認(rèn) ReplaceNodeSubsetsWithDelegateKernels 的返回狀態(tài)。如果返回kTfLiteError后續(xù)就不會(huì)有 delegate kernel 節(jié)點(diǎn)。確認(rèn) delegate kernel 的 invoke 真的被調(diào)了。在最外層 invoke 打日志如果沒(méi)日志說(shuō)明執(zhí)行計(jì)劃里根本不存在你的 delegate kernel。確認(rèn)沒(méi)有其他 delegate 搶先接管了同一批節(jié)點(diǎn)。多個(gè) delegate 疊加時(shí)先執(zhí)行的 delegate 可能已經(jīng)把 ADD 節(jié)點(diǎn)替換掉了后注冊(cè)的 delegate 自然匹配不到。這套順序看起來(lái)簡(jiǎn)單但執(zhí)行的時(shí)候一定要借助日志而不是靠“我感覺(jué)”。TFLite 編譯時(shí)如果開(kāi)了 verbose log會(huì)在ModifyGraphWithDelegate階段打印節(jié)點(diǎn)替換的詳細(xì)信息沒(méi)開(kāi)日志時(shí)自己在 Prepare 和 kernel 函數(shù)里埋 printf 是最快的。另一個(gè)有價(jià)值的經(jīng)驗(yàn)是delegate 接管率高不等于端到端延遲一定更低。如果被接管節(jié)點(diǎn)夾在大量 CPU 節(jié)點(diǎn)之間tensor 數(shù)據(jù)反復(fù)在 CPU 和硬件后端之間拷貝開(kāi)銷(xiāo)可能抵消掉加速收益。所以做 delegate 優(yōu)化時(shí)我會(huì)把執(zhí)行計(jì)劃畫(huà)出來(lái)重點(diǎn)看能形成多大塊的連續(xù)子圖而不是追求接管的節(jié)點(diǎn)數(shù)量。這也是為什么前面強(qiáng)調(diào)ReplaceNodeSubsetsWithDelegateKernels是按子圖接管——它給了 delegate 做整塊調(diào)度的機(jī)會(huì)而這個(gè)機(jī)會(huì)需要你在 Prepare 里主動(dòng)用好。最后說(shuō)一句我自己的感受。注冊(cè)機(jī)制看懂了之后TFLite 的很多“玄學(xué)”問(wèn)題會(huì)變得特別直白模型文件里的算子描述是一回事解釋器里的內(nèi)核注冊(cè)是另一回事delegate 又是在執(zhí)行計(jì)劃層面的第三回事。三層之間通過(guò)FindOp和ReplaceNodeSubsetsWithDelegateKernels這兩個(gè)關(guān)鍵點(diǎn)串聯(lián)起來(lái)。以后再遇到“明明注冊(cè)了為什么沒(méi)生效”“delegate 接了為什么沒(méi)加速”這類(lèi)問(wèn)題順著報(bào)錯(cuò)信息回到這一層層的鏈路里去定位多半就會(huì)豁然開(kāi)朗。