uTensor: 轻量级边缘 AI 推理框架

uTensor 简介

uTensor 是一个极其轻量级的机器学习推理框架,专为 Arm 目标设备优化。它基于 TensorFlow 构建,旨在将深度学习模型部署到资源受限的嵌入式设备上。uTensor 由两个主要部分组成:运行时库和离线工具。运行时库提供了推理所需的核心功能,而离线工具则负责处理大部分模型转换工作。

uTensor 的主要特点

极致轻量级:核心运行时库仅约 2KB 大小
针对 Arm 设备优化:专为 Arm Cortex-M 系列等微控制器设计
基于 TensorFlow:可以直接使用 TensorFlow 训练的模型
简单易用:提供高层 API,使用简单直观
灵活可扩展:支持自定义算子、内存管理器等

uTensor 的设计理念是"小而美",致力于在有限的资源下实现高效的AI推理。它非常适合IoT设备、可穿戴设备等对功耗和内存有严格要求的场景。

uTensor 工作流程

uTensor 的工作流程主要包括以下步骤:

使用 TensorFlow 构建并训练模型
使用 uTensor 离线工具将模型转换为 C++ 代码
将生成的代码集成到嵌入式项目中
在目标设备上编译并运行推理代码

uTensor workflow

如上图所示,开发者首先在 TensorFlow 中构建和训练模型。然后,uTensor 的离线工具会将模型转换为 .cpp 和 .hpp 文件,这些文件包含了推理所需的 C++11 代码。最后,开发者只需将这些生成的代码集成到嵌入式项目中即可。这种方法大大简化了将深度学习模型部署到嵌入式设备的过程。

uTensor 核心设计理念

uTensor 的设计围绕以下几个关键理念:

张量(Tensor)描述数据的访问方式和来源
- 操作符的性能取决于使用的张量类型
操作符(Operator)与张量无关
- 高性能操作符可以一次获取大块数据
追求低总功耗
低静态和动态内存占用
- 整个系统中每个张量的成本都很低
- 轻量级的类层次结构

此外,uTensor 还特别注重以下几点:

系统安全性
- 所有张量元数据和实际数据都存储在专用区域
- 可以在代码生成时或编译时保证运行时使用的 RAM 不超过指定值
- 通常不应与用户空间或系统空间内存发生冲突
清晰、简洁和可调试性
- 开发者可以在不依赖代码生成的情况下修改模型
- 开发者可以轻松理解模型文件的图结构
- 默认张量接口行为类似高级语言,但利用 C++ 的速度
- 可扩展、可配置和可优化的错误处理

这些设计理念不仅体现在代码中,还反映在代码结构上,使得 uTensor 具有良好的可维护性、可扩展性和用户友好性。

uTensor 高层 API 示例

以下是一个使用 uTensor 高层 API 的简单示例:

using namespace uTensor;

const uint8_t s_a[4] = {1, 2, 3, 4};
const uint8_t s_b[4] = {5, 6, 7, 8};
const uint8_t s_c_ref[4] = {19, 22, 43, 50};

// 创建内存分配器
localCircularArenaAllocator<256> meta_allocator;
localCircularArenaAllocator<256> ram_allocator;

void foo() {
  // 设置 uTensor 上下文使用的分配器
  Context::get_default_context()->set_metadata_allocator(&meta_allocator);
  Context::get_default_context()->set_ram_data_allocator(&ram_allocator);

  // 创建张量
  Tensor a = new RomTensor({2, 2}, u8, s_a);
  Tensor b = new RomTensor({2, 2}, u8, s_b);
  Tensor c_ref = new RomTensor({2,2}, u8, s_c_ref);
  Tensor c = new RamTensor({2, 2}, u8);

  // 创建并执行矩阵乘法操作
  MatrixMultOperator<uint8_t> mult_AB;
  mult_AB
      .set_inputs({{MatrixMultOperator<uint8_t>::a, a}, {MatrixMultOperator<uint8_t>::b, b}})
      .set_outputs({{MatrixMultOperator<uint8_t>::c, c}})
      .eval();

  // 比较结果
  TensorShape& c_shape = c->get_shape();
  for (int i = 0; i < c_shape[0]; i++) {
    for (int j = 0; j < c_shape[1]; j++) {
      if( static_cast<uint8_t>(c(i, j)) != static_cast<uint8_t>(c_ref(i, j)) ) {
        printf("Oh crap!\n");
        exit(-1);
      }
    }
  }
}

这个例子展示了如何使用 uTensor 创建张量、执行矩阵乘法操作并验证结果。它突出了 uTensor API 的简洁性和易用性。

在不同平台上构建和运行 uTensor

本地构建和测试

要在本地构建和测试 uTensor,可以按照以下步骤操作:

git clone git@github.com:uTensor/uTensor.git
cd uTensor/
git checkout proposal/rearch
git submodule init
git submodule update
mkdir build
cd build/
cmake -DPACKAGE_TESTS=ON -DCMAKE_BUILD_TYPE=Debug ..
make
make test

在 Arm Mbed OS 上构建和运行

uTensor 核心库已配置为即插即用的 Mbed 库。要在 Mbed OS 项目中使用它,只需按以下方式导入:

mbed new my_project
cd my_project
mbed import https://github.com/uTensor/uTensor.git
# 创建主文件
# 运行 uTensor-cli 工作流并复制模型目录到此处
mbed compile # 正常编译

在 Arm 系统上构建和运行

对于 Arm 系统,uTensor 提供了两种构建方式:

默认构建:

mkdir build && cd build
cmake -DCMAKE_BUILD_TYPE=Debug -DCMAKE_TOOLCHAIN_FILE=../extern/CMSIS_5/CMSIS/DSP/gcc.cmake  ..

使用 CMSIS 优化内核:

mkdir build && cd build
cmake -DARM_PROJECT=1 -DCMAKE_BUILD_TYPE=Debug -DCMAKE_TOOLCHAIN_FILE=../extern/CMSIS_5/CMSIS/DSP/gcc.cmake  ..

uTensor 的应用场景

uTensor 主要适用于以下场景:

物联网(IoT)设备:利用 uTensor 可以在资源受限的 IoT 设备上实现本地AI推理,提高响应速度和隐私保护。
可穿戴设备:uTensor 的轻量级特性使其非常适合部署在智能手表、健康监测设备等可穿戴设备上。
智能家居:在智能音箱、智能摄像头等家居设备中集成 uTensor,可以实现更智能的本地控制和决策。
工业物联网:在工业传感器和控制器中使用 uTensor,可以实现更智能的实时监控和预测性维护。
自动驾驶:在车载系统中集成 uTensor,可以实现一些轻量级的实时AI决策任务。
机器人:在小型机器人或无人机中使用 uTensor,可以实现本地的图像识别、语音处理等AI功能。

uTensor 的未来发展

作为一个活跃的开源项目,uTensor 正在不断发展和改进。未来的发展方向可能包括:

支持更多类型的神经网络模型和层
进一步优化推理性能和内存使用
提供更多预训练模型和示例应用
增强与其他嵌入式开发框架的集成
改进开发工具和调试功能
扩展到更多的硬件平台

结语

uTensor 为将深度学习模型部署到资源受限的嵌入式设备提供了一个强大而灵活的解决方案。它的轻量级设计、优化的性能和易用的 API 使其成为边缘 AI 和 TinyML 领域的重要工具。随着物联网和边缘计算的快速发展,uTensor 有望在未来发挥更加重要的作用,推动 AI 技术向更广泛的嵌入式设备普及。

对于希望在资源受限设备上实现 AI 功能的开发者来说,uTensor 无疑是一个值得关注和尝试的框架。它不仅能够帮助开发者快速将 TensorFlow 模型部署到嵌入式设备,还能确保在有限的资源下实现高效的推理。随着项目的不断发展和社区的持续贡献,我们可以期待 uTensor 在未来为嵌入式 AI 领域带来更多创新和可能性。