:從ONNX模型到CUBE-AI推理全流程)
1. 為什么要在STM32上跑AI模型1.1 從“云端推理”到“邊緣推理”的轉(zhuǎn)變過去幾年做AI應(yīng)用的主流思路是把數(shù)據(jù)傳到云端服務(wù)器在GPU集群上跑推理再把結(jié)果返回到終端設(shè)備。這套模式在帶寬充足、延遲不敏感的場景下沒問題但一旦落到工業(yè)現(xiàn)場、消費電子、車載設(shè)備這些領(lǐng)域麻煩就來了網(wǎng)絡(luò)不穩(wěn)定、數(shù)據(jù)隱私敏感、響應(yīng)延遲要求高、長期流量成本壓不住。于是“邊緣AI”這個概念被反復(fù)提起而STM32作為嵌入式領(lǐng)域出貨量最大的MCU家族之一自然成了很多人嘗試部署輕量級AI模型的首選平臺。我最早接觸這個方向是因為一個電機異常檢測的項目??蛻粢笤诓宦?lián)網(wǎng)的前提下用振動傳感器采集數(shù)據(jù)實時判斷電機運行狀態(tài)。一開始想用樹莓派但成本、功耗、供貨周期都不合適最后選了一顆STM32F407配合CUBE-AI把一個小型神經(jīng)網(wǎng)絡(luò)塞進去整個方案BOM成本壓到了原來方案的三分之一。從那以后我陸續(xù)在STM32F4、F7、H7、L4、U5等多個系列上部署過模型踩了不少坑也總結(jié)了一些相對成熟的流程。這篇文章面向的是有一定STM32開發(fā)基礎(chǔ)、想嘗試AI模型部署但不知道從哪下手的工程師。我會把整個流程拆開從模型訓(xùn)練、格式轉(zhuǎn)換、量化、CUBE-AI集成到最終在板子上跑起來每一步都給出可復(fù)現(xiàn)的操作和參數(shù)說明。10分鐘是理想情況下的時間前提是你已經(jīng)有一個訓(xùn)練好的模型和一塊能用的開發(fā)板。1.2 STM32做AI推理的硬件底子不是所有STM32都能跑AI模型選型的時候要看幾個關(guān)鍵指標。首先是Flash和RAM模型權(quán)重和中間激活值都要占空間一個幾十KB的模型在F103上基本沒戲但在F407或H743上就很輕松。其次是主頻和是否有FPU浮點運算單元帶FPU的芯片做浮點推理速度會快很多比如F4系列有單精度FPUH7系列有雙精度FPU而F0、F1這些老型號沒有FPU跑浮點模型會非常吃力。再就是是否有DSP指令集Cortex-M4和M7都支持DSP擴展做卷積、矩陣乘法這類操作時能明顯加速。下面這張表是我實際用過的幾款芯片在跑同一個關(guān)鍵詞識別模型時的表現(xiàn)模型輸入是40x10的MFCC特征輸出是12類關(guān)鍵詞量化到int8供你選型時參考芯片型號主頻FPUDSPFlash占用RAM占用單次推理耗時STM32F103C872MHz無無38KB12KB約420msSTM32F407VG168MHz單精度有36KB10KB約28msSTM32F767ZI216MHz雙精度有36KB10KB約18msSTM32H743ZI480MHz雙精度有36KB10KB約9msSTM32L4R5ZI120MHz單精度有36KB10KB約45ms從表里能看出來F103雖然能跑但420ms的延遲在很多實時場景下是不可接受的。F407是性價比拐點28ms已經(jīng)能滿足大部分中低速控制場景。H743適合對延遲要求苛刻的應(yīng)用比如音頻實時處理。L4系列主打低功耗適合電池供電的場合但速度會慢一些。注意上表中的Flash和RAM占用是CUBE-AI生成代碼后的實際編譯結(jié)果不同版本的CUBE-AI和不同的優(yōu)化等級會有差異建議以你實際工程編譯后的map文件為準。2. 模型準備從訓(xùn)練到ONNX2.1 模型選型與訓(xùn)練框架STM32上能跑的模型有幾個硬約束參數(shù)量不能太大一般控制在100KB以內(nèi)比較穩(wěn)妥輸入維度不能太高比如圖像輸入224x224x3對MCU來說就太重了通常要降到32x32或48x48網(wǎng)絡(luò)結(jié)構(gòu)要簡單避免復(fù)雜的自定義算子因為CUBE-AI對算子的支持是有限的。我常用的訓(xùn)練框架是TensorFlow/Keras和PyTorch。Keras的好處是API簡潔導(dǎo)出ONNX或直接導(dǎo)出SavedModel都方便。PyTorch在學(xué)術(shù)界更流行導(dǎo)出ONNX的流程也很成熟。不管用哪個框架最終目標都是得到一個ONNX格式的模型文件因為CUBE-AI對ONNX的支持最完善。以一個簡單的振動分類模型為例輸入是三軸加速度計的256點采樣輸出是5種狀態(tài)正常、不平衡、松動、軸承故障、齒輪故障。網(wǎng)絡(luò)結(jié)構(gòu)用三層一維卷積加兩層全連接參數(shù)量控制在20KB左右。訓(xùn)練的時候要注意輸入數(shù)據(jù)要做歸一化把原始加速度值映射到[-1, 1]或[0, 1]區(qū)間這樣量化后的精度損失會小很多。2.2 導(dǎo)出ONNX的實操細節(jié)從PyTorch導(dǎo)出ONNX核心代碼就幾行import torch import torch.onnx # 假設(shè)model是你的網(wǎng)絡(luò)input是示例輸入 model.eval() dummy_input torch.randn(1, 3, 256) # batch1, 三軸, 256點 torch.onnx.export( model, dummy_input, vibration_model.onnx, input_names[accel_input], output_names[state_output], opset_version11, dynamic_axesNone # MCU上batch固定為1不需要動態(tài)軸 )這里有幾個關(guān)鍵點。opset_version建議用11或12太新的版本CUBE-AI可能不支持太老的版本某些算子導(dǎo)出會有問題。dynamic_axes一定要設(shè)為None或者不設(shè)置因為MCU上batch size固定為1動態(tài)軸會引入額外的復(fù)雜度。model.eval()必須調(diào)用否則BatchNorm和Dropout層的行為會和推理時不一致。導(dǎo)出完成后用onnxruntime驗證一下模型是否能正常推理輸入輸出shape是否符合預(yù)期import onnxruntime as ort import numpy as np sess ort.InferenceSession(vibration_model.onnx) input_name sess.get_inputs()[0].name output_name sess.get_outputs()[0].name test_input np.random.randn(1, 3, 256).astype(np.float32) result sess.run([output_name], {input_name: test_input}) print(result[0].shape) # 應(yīng)該是(1, 5)如果這一步報錯說明ONNX模型本身有問題先解決再往下走不要帶著問題模型去跑CUBE-AI否則后面排查起來更麻煩。2.3 量化int8到底怎么選量化是MCU部署AI模型的關(guān)鍵步驟。STM32的CUBE-AI支持float32和int8兩種主要的數(shù)據(jù)類型。float32精度高但占用大、速度慢int8占用小、速度快但會有精度損失。實際項目中我?guī)缀蹩偸莾?yōu)先考慮int8只有在精度實在壓不住的情況下才退回float32。CUBE-AI的量化是訓(xùn)練后量化Post-Training Quantization不需要重新訓(xùn)練模型。它需要一個校準數(shù)據(jù)集通常從訓(xùn)練集里隨機抽100到500個樣本就夠了。校準數(shù)據(jù)要覆蓋各種工況比如振動模型里要包含正常和各種故障狀態(tài)的樣本否則量化后的模型在某些類別上會嚴重偏移。量化的核心參數(shù)是quantize選項在CUBE-AI的配置里可以選int8、float32或者int8_float32混合模式?;旌夏J绞侵覆糠謱佑胕nt8部分層用float32適合那些對精度特別敏感的層。我一般先用純int8跑一遍看精度下降多少如果下降在2%以內(nèi)就接受超過5%就考慮混合模式或者調(diào)整網(wǎng)絡(luò)結(jié)構(gòu)。實操心得量化后的模型精度驗證一定要用獨立的測試集不要用校準集本身。我見過有人用校準集驗證精度看起來很好但實際部署后效果很差就是因為校準集過擬合了。3. CUBE-AI集成從ONNX到C代碼3.1 STM32CubeMX和CUBE-AI的安裝配置CUBE-AI是ST官方推出的工具集成在STM32CubeMX里。如果你用的是STM32CubeIDE它也內(nèi)置了CUBE-AI插件。我習慣用CubeMX獨立版因為版本更新更及時而且可以單獨管理CUBE-AI的版本。安裝步驟大致是先裝STM32CubeMX然后在Help菜單里選Manage embedded software packages找到X-CUBE-AI這個包勾選你需要的版本。截至我寫這篇文章的時候X-CUBE-AI已經(jīng)更新到8.x版本對ONNX的算子支持比早期版本好了很多。安裝完成后新建工程或者打開已有工程在Middleware里就能看到X-CUBE-AI的選項。這里有個坑要注意CUBE-AI的版本和CubeMX的版本有對應(yīng)關(guān)系不是隨便組合都能用。比如CubeMX 6.8配CUBE-AI 8.0是穩(wěn)定的但配7.x可能會報錯。建議直接用CubeMX推薦的默認版本不要強行混搭。3.2 添加模型和配置參數(shù)在CubeMX里啟用X-CUBE-AI后界面會多出一個AI選項卡。點擊Add network選擇你導(dǎo)出的ONNX文件。CUBE-AI會自動分析模型結(jié)構(gòu)顯示每一層的類型、輸入輸出shape、參數(shù)量等信息。這時候要重點看幾個地方第一有沒有不支持的算子。如果某個層顯示為紅色或者有警告說明CUBE-AI不支持這個算子需要回到模型層面替換成支持的算子。常見的坑包括自定義的激活函數(shù)、非標準的池化方式、復(fù)雜的reshape操作等。第二輸入輸出的數(shù)據(jù)類型和shape。CUBE-AI默認會把輸入設(shè)為float32如果你要用int8推理需要在配置里把輸入類型改成int8同時注意輸入數(shù)據(jù)的歸一化方式要和訓(xùn)練時一致。第三內(nèi)存分配。CUBE-AI會給出一個預(yù)估的Flash和RAM占用你可以調(diào)整memory pool的大小。如果RAM不夠可以嘗試開啟use activation buffer選項讓中間激活值復(fù)用同一塊內(nèi)存代價是推理速度會稍微慢一點。配置完成后點擊Generate CodeCUBE-AI會在工程里生成一系列文件包括網(wǎng)絡(luò)權(quán)重、推理引擎、API接口等。生成的文件通常在Middlewares/ST/AI目錄下核心的API在ai_platform.h和app_x-cube-ai.h里。3.3 生成代碼的結(jié)構(gòu)解析CUBE-AI生成的代碼結(jié)構(gòu)比較清晰主要分三部分。第一部分是模型定義在network.c和network.h里包含了每一層的權(quán)重數(shù)據(jù)和網(wǎng)絡(luò)拓撲。第二部分是推理引擎在ai_platform.c里負責調(diào)度各層的計算。第三部分是應(yīng)用接口在app_x-cube-ai.c里提供了初始化、推理、獲取結(jié)果的函數(shù)。關(guān)鍵API有三個// 初始化AI模型 ai_handle ai_model_init(void); // 執(zhí)行推理 ai_error ai_model_run(ai_handle handle, const ai_buffer* input, ai_buffer* output); // 獲取推理結(jié)果 ai_buffer* ai_model_get_output(ai_handle handle);實際使用的時候你需要在main函數(shù)里先調(diào)用初始化然后在主循環(huán)或定時器中斷里填充輸入buffer調(diào)用推理函數(shù)最后讀取輸出buffer。輸入buffer的填充要注意數(shù)據(jù)格式如果是int8量化模型輸入值要按量化參數(shù)縮放成int8整數(shù)。4. 在STM32上跑通第一個推理4.1 工程搭建與編譯假設(shè)你已經(jīng)用CubeMX生成了一個帶CUBE-AI的工程接下來要做的就是把工程導(dǎo)入到你的IDE里編譯。我用的是STM32CubeIDE因為它是免費的而且和CubeMX無縫集成。如果你習慣用Keil或IAR也可以但要注意CUBE-AI生成的代碼對編譯器的C標準有要求建議用C11或更高。編譯之前檢查幾個地方一是堆棧大小AI推理會用到比較大的??臻g建議把主棧調(diào)到至少4KB堆調(diào)到至少8KB二是優(yōu)化等級建議用-O2或-O3-O0會讓推理速度慢好幾倍三是浮點ABI如果芯片有FPU要在編譯選項里開啟hard float否則浮點運算會用軟件模擬速度差一個數(shù)量級。編譯通過后先別急著跑推理用調(diào)試器看一下生成的模型信息。CUBE-AI提供了一個ai_model_info()函數(shù)可以打印出模型的輸入輸出shape、參數(shù)量、量化參數(shù)等。把這些信息和你在PC上驗證的ONNX模型對比一下確認一致后再往下走。4.2 輸入數(shù)據(jù)的預(yù)處理MCU上的輸入數(shù)據(jù)通常來自傳感器比如加速度計、麥克風、攝像頭。這些原始數(shù)據(jù)不能直接喂給模型需要做預(yù)處理。以振動模型為例加速度計輸出的是三軸16位整數(shù)要先轉(zhuǎn)成浮點再做歸一化最后按模型要求的shape排列。預(yù)處理這一步很容易出問題。我遇到過好幾次推理結(jié)果完全不對最后發(fā)現(xiàn)是輸入數(shù)據(jù)的排列順序和訓(xùn)練時不一致。比如訓(xùn)練時數(shù)據(jù)是[通道][采樣點]的格式部署時不小心寫成了[采樣點][通道]模型完全懵了。解決辦法是在PC上用同樣的預(yù)處理代碼跑一遍把結(jié)果和MCU上的結(jié)果對比確認一致后再繼續(xù)。如果模型是int8量化的輸入數(shù)據(jù)還要做量化。CUBE-AI會給出輸入的scale和zero_point參數(shù)你需要按這個公式把浮點值轉(zhuǎn)成int8int8_value round(float_value / scale) zero_point這個計算可以在PC上預(yù)先算好也可以放在MCU上實時算。如果傳感器數(shù)據(jù)變化不快建議在MCU上實時算靈活性更好。4.3 推理執(zhí)行與結(jié)果解析推理函數(shù)的調(diào)用比較簡單但有幾個細節(jié)要注意。第一輸入buffer和輸出buffer要用CUBE-AI提供的ai_buffer結(jié)構(gòu)體不要自己隨便定義數(shù)組。第二推理函數(shù)是阻塞的調(diào)用后會一直等到計算完成才返回所以在實時性要求高的場景下要考慮把推理放在低優(yōu)先級任務(wù)里或者用DMA加中斷的方式做數(shù)據(jù)搬運。推理完成后輸出buffer里是量化后的int8值需要反量化回浮點才能得到有意義的分類概率。反量化公式是float_value (int8_value - zero_point) * scale如果是分類任務(wù)通常還要做softmax把輸出轉(zhuǎn)成概率。但注意CUBE-AI生成的模型如果最后一層是softmax輸出已經(jīng)是概率了不需要再做一次。如果不確定可以看模型結(jié)構(gòu)里最后一層的類型。避坑技巧第一次跑推理的時候建議先用一組已知結(jié)果的數(shù)據(jù)做驗證。比如從測試集里拿一個樣本在PC上跑一遍記錄輸出然后在MCU上跑同樣的樣本對比兩者的輸出差異。如果差異在合理范圍內(nèi)比如最大絕對誤差小于0.05說明部署成功。如果差異很大就要逐步排查是預(yù)處理、量化還是推理引擎的問題。5. 性能優(yōu)化與內(nèi)存壓縮5.1 推理速度的優(yōu)化手段推理速度是MCU部署AI模型時最常被問到的問題。優(yōu)化手段分幾個層面。最直接的是提高主頻把芯片跑到最高頻率比如F407從168MHz超到180MHz雖然官方不推薦長期超頻H743從480MHz跑到550MHz。但超頻有風險量產(chǎn)項目不建議。第二個層面是優(yōu)化模型結(jié)構(gòu)。減少層數(shù)、減小通道數(shù)、用深度可分離卷積替代標準卷積這些都能顯著降低計算量。我做過一個對比把標準卷積換成深度可分離卷積后參數(shù)量降到原來的三分之一推理速度提升了一倍多精度只掉了1.5%。第三個層面是用CUBE-AI的優(yōu)化選項。CUBE-AI在生成代碼時有幾個優(yōu)化等級比如balanced、time、ram。選time會生成更快的代碼但占用更多Flash選ram會壓縮內(nèi)存占用但速度慢一些。根據(jù)你的實際約束來選。第四個層面是用CMSIS-NN庫。CUBE-AI生成的代碼底層會調(diào)用CMSIS-NN的優(yōu)化函數(shù)但前提是你的芯片支持DSP指令集。如果芯片是Cortex-M4或M7確保在工程里啟用了CMSIS-NN否則會退回到普通的C實現(xiàn)速度差很多。5.2 內(nèi)存占用的壓縮策略內(nèi)存是另一個瓶頸。STM32F407有192KB RAM聽起來不少但模型權(quán)重、中間激活值、輸入輸出buffer、再加上你的應(yīng)用程序很容易就超了。壓縮內(nèi)存的策略有幾個一是用int8量化權(quán)重和激活值都從4字節(jié)降到1字節(jié)直接省75%的內(nèi)存。二是開啟activation buffer復(fù)用讓不同層的中間結(jié)果共用同一塊內(nèi)存CUBE-AI會自動計算最優(yōu)的內(nèi)存復(fù)用方案。三是把模型權(quán)重放到外部Flash運行時按需加載但這會增加推理延遲適合對速度不敏感的場景。還有一個容易被忽略的點是堆棧大小。AI推理的調(diào)用??赡芎苌钐貏e是層數(shù)多的模型。如果棧太小會出現(xiàn)莫名其妙的hard fault。我一般會把主棧設(shè)成8KB堆設(shè)成16KB然后根據(jù)實際使用情況調(diào)整。5.3 精度與速度的權(quán)衡精度和速度永遠是一對矛盾。我的經(jīng)驗是先確定你能接受的最低精度然后在這個約束下盡量優(yōu)化速度。比如振動分類模型如果客戶要求準確率不低于95%那量化后的模型必須達到這個線達不到就退回float32或者調(diào)整網(wǎng)絡(luò)結(jié)構(gòu)。量化精度損失主要來自兩個方面權(quán)重的不均勻分布和激活值的動態(tài)范圍。如果某一層的權(quán)重集中在很小的范圍內(nèi)量化后很多值會變成0信息就丟了。解決辦法是在訓(xùn)練時加入量化感知訓(xùn)練QAT讓模型提前適應(yīng)量化誤差。不過QAT需要重新訓(xùn)練流程更復(fù)雜適合對精度要求極高的場景。6. 常見問題與排查實錄6.1 推理結(jié)果完全不對這是最常見的問題原因通常有幾個。第一輸入數(shù)據(jù)格式不對比如通道順序、歸一化方式、量化參數(shù)和訓(xùn)練時不一致。第二模型導(dǎo)出ONNX時出了問題比如opset版本不匹配、某些層被錯誤簡化。第三CUBE-AI版本和模型不兼容某些算子被錯誤處理。排查方法是從PC端開始逐級對比。先在PC上用ONNX Runtime跑一遍記錄輸出。然后在PC上用CUBE-AI的模擬器跑一遍CUBE-AI提供了PC端的驗證工具對比輸出。如果這兩步一致說明模型和CUBE-AI都沒問題問題出在MCU端的預(yù)處理或數(shù)據(jù)搬運。如果這兩步不一致說明ONNX導(dǎo)出或CUBE-AI配置有問題。6.2 編譯報錯和鏈接錯誤CUBE-AI生成的代碼對編譯器和鏈接器有一些要求。常見的報錯包括找不到ai_platform.h說明頭文件路徑?jīng)]加對undefined reference to ai_model_init說明生成的源文件沒加入編譯region RAM overflowed說明內(nèi)存不夠需要調(diào)整鏈接腳本或壓縮模型。鏈接腳本的問題在STM32上很常見特別是用CubeMX生成的工程默認的鏈接腳本可能沒有給AI模型留足夠的空間。你需要手動修改.ld文件把AI相關(guān)的段放到合適的位置。如果用的是H7系列還要注意DTCM、AXI SRAM、SRAM1/2/3的分配不同內(nèi)存區(qū)域的訪問速度不一樣把權(quán)重放在AXI SRAM、激活值放在DTCM通常是最優(yōu)的。6.3 推理速度比預(yù)期慢如果推理速度比預(yù)期慢很多先檢查編譯優(yōu)化等級。我見過有人用-O0編譯推理耗時是-O2的5倍。然后檢查FPU是否啟用沒有FPU的話浮點運算會慢得離譜。再檢查CMSIS-NN是否啟用沒有DSP指令集加速的話卷積運算會慢很多。還有一個隱藏的坑是中斷干擾。如果推理過程中頻繁被高優(yōu)先級中斷打斷實際耗時會被拉長。解決辦法是把推理放在臨界區(qū)里或者用DMA把數(shù)據(jù)搬運和推理并行起來。6.4 常見問題速查表現(xiàn)象可能原因排查方法解決方案推理結(jié)果全為同一類輸入數(shù)據(jù)未歸一化或量化參數(shù)錯誤對比PC和MCU的輸入buffer檢查歸一化和量化公式編譯報錯找不到AI頭文件頭文件路徑未添加查看編譯輸出在IDE里添加Middlewares/ST/AI路徑鏈接報錯RAM溢出模型太大或堆棧設(shè)置過小查看map文件壓縮模型或調(diào)整鏈接腳本推理速度極慢優(yōu)化等級為-O0或FPU未啟用查看編譯選項改為-O2并啟用hard float運行中hard fault棧溢出或內(nèi)存越界用調(diào)試器查看fault寄存器增大??臻g檢查buffer大小量化后精度驟降校準集不具代表性用測試集驗證更換校準集或改用混合量化最后分享一個我踩過的坑有一次模型在F407上跑得好好的換到F103上就完全不對。排查了半天發(fā)現(xiàn)是F103沒有FPUCUBE-AI生成的代碼默認用浮點在F103上浮點運算被軟件模擬不僅慢而且某些中間結(jié)果因為精度問題出現(xiàn)了溢出。解決辦法是在CUBE-AI配置里把數(shù)據(jù)類型改成純int8避免任何浮點運算。所以換芯片的時候一定要重新檢查CUBE-AI的配置不要直接復(fù)制工程。