2K內(nèi)存單片機就能跑的嵌入式AI模型

這兩年隨著ChatGPT、DeepSeek的火爆，AI已經(jīng)遍布工作和生活的各個角落，嵌入式端側AI也逐漸發(fā)展起來了。

今天就來分享一個可用于2KB內(nèi)存單片機的嵌入式AI模型：uTensor。

關于?uTensor 模型

uTensor 是一個基于 Tensorflow 構建的極其輕量級的機器學習推理框架，并針對 Arm 處理器進行了優(yōu)化。它由一個運行時庫和一個處理大部分模型轉換工作的離線工具組成。

模型地址：https://github.com/uTensor/uTensor

此存儲庫包含核心運行時和運算符、內(nèi)存管理器、調(diào)度器等的一些示例實現(xiàn)，核心運行時的大小僅為：2KB！

uTensor只需要2KB內(nèi)存的輕量化設計特點，就是實現(xiàn)了極致壓縮：將TensorFlow模型轉換為.cpp、.hpp源代碼，消除冗余依賴。同時，預分配內(nèi)存區(qū)域，杜絕運行時內(nèi)存的泄漏。

實測核心運行時和基礎算子的總代碼量僅2KB，相當于一張圖片的1/1000.

uTensor 工作原理

uTensor 工作原理大致如下圖所示：

在 Tensorflow 中構建和訓練模型，uTensor 獲取模型并生成 .cpp 和 .hpp 源文件。這些文件包含生成的推理所需的 C++代碼，只需要把生成的源文件復制到你的嵌入式項目中即可，實現(xiàn)過程非常簡單。

uTensor 運行時由兩個主要組件組成：

uTensor Core：其中包含滿足 uTensor 性能運行時契約所需的基本數(shù)據(jù)結構、接口和類型等。

uTensor 庫：作為一系列基于 uTensor Core 構建的默認實現(xiàn)。

構建系統(tǒng)分別編譯這兩個組件，使用戶能夠輕松擴展和覆蓋構建在 uTensor 核心之上的實現(xiàn)，例如自定義內(nèi)存管理器、張量、運算符和錯誤處理程序。

錯誤處理程序：

SimpleErrorHandler errH(50); // Maintain a history of 50 eventsContext::get_default_context()->set_ErrorHandler(&errH);...// A bunch of allocations...
// Check to make sure a rebalance has occurred inside our allocatorbool has_rebalanced = std::find(errH.begin(), errH.end(), localCircularArenaAllocatorRebalancingEvent()) != errH.end();

Tensor 讀寫接口：

uint8_t myBuffer[4] = { 0xde, 0xad, 0xbe, 0xef };Tensor mTensor = new BufferTensor({2,2}, u8, myBuffer); // define a 2x2 tensor of uint8_ts
uint8_t a1 = mTensor(0,0);  // implicitly casts the memory referenced at this index to a uint8_tprintf("0x%hhxn", a1);     // prints 0xde
uint16_t a2 = mTensor(0,0); // implicitly casts the memory referenced at this index to a uint16_tprintf("0x%hxn", a2);      // prints 0xdead
uint32_t a3 = mTensor(0,0); // implicitly casts the memory referenced at this index to a uint32_tprintf("0x%xn", a3);      // prints 0xdeadbeef
// You can also write and read values with explicit casting and get similar behaviormTensor(0,0) = static_cast<uint8_t>(0xFF);printf("0xhhxn", static_cast<uint8_t>(mTensor(0,0)));

出于性能原因，各種 Tensor 讀/寫接口更像緩沖區(qū)，而不是成熟的 C++ 類型化對象，盡管高級接口本質上看起來非常 Pythonic 。實際的讀取和寫入取決于用戶如何轉換此緩沖區(qū)。

uTensor 構建、運行和測試

官方給出了 uTensor 構建、運行和測試的一些方法。

比如在本地構建和測試：

git clone git@github.com:uTensor/uTensor.gitcd uTensor/git checkout proposal/rearchgit submodule initgit submodule updatemkdir buildcd build/cmake -DPACKAGE_TESTS=ON -DCMAKE_BUILD_TYPE=Debug ..makemake test

在 Arm Mbed OS 上構建和運行：

mbed new my_projectcd my_projectmbed import https://github.com/uTensor/uTensor.git# Create main file# Run uTensor-cli workflow and copy model directory herembed compile # as normal

還有在在Arm 系統(tǒng)上構建和運行：

mkdir build && cd buildcmake -DCMAKE_BUILD_TYPE=Debug -DCMAKE_TOOLCHAIN_FILE=../extern/CMSIS_5/CMSIS/DSP/gcc.cmake  ..
//使用?CMSIS?優(yōu)化內(nèi)核mkdir build && cd buildcmake -DARM_PROJECT=1 -DCMAKE_BUILD_TYPE=Debug -DCMAKE_TOOLCHAIN_FILE=../extern/CMSIS_5/CMSIS/DSP/gcc.cmake  ..

以上只是提供了一些參考和思路，實現(xiàn)的具體細節(jié)，需要大家進一步結合 uTensor 模型進行優(yōu)化。