TinyML - TVM 如何驾驭微型设备

2020年6月4日 • Logan Weber 和 Andrew Reusch，OctoML

低成本、人工智能驱动的消费设备的普及，使得机器学习研究人员和从业者对“裸机”（低功耗，通常没有操作系统）设备产生了广泛的兴趣。虽然专家已经可以在某些裸机设备上运行一些模型，但针对各种设备优化模型仍然具有挑战性，通常需要手动优化的设备特定库。对于那些例如没有 Linux 支持的平台，目前还没有可扩展的解决方案来部署模型。因此，为了针对新设备，开发人员必须实现一次性的自定义软件堆栈来管理系统资源和调度模型执行。

机器学习软件的手动优化并非裸机设备领域独有。事实上，这一直是开发人员使用其他硬件后端（例如，GPU 和 FPGA）时面临的常见问题。TVM 已被证明能够适应不断涌现的新硬件目标，但到目前为止，它还无法应对微控制器的独特特性。为了解决该领域的问题，我们扩展了 TVM，使其具有微控制器后端，称为 µTVM（脚注：发音为“MicroTVM”）。µTVM 促进了张量程序在裸机设备上的主机驱动执行，并通过 AutoTVM（TVM 内置的张量程序优化器）实现了这些程序的自动优化。下图显示了 µTVM + AutoTVM 基础设施的概览。

/images/microtvm/autotvm-infrastructure.png

让我们看看实际应用

在我们讨论 TVM/MicroTVM 是什么或它是如何工作之前，让我们先看一个它在实际应用中的快速示例。

/images/microtvm/hardware-connection-diagram.png
一个标准的 µTVM 设置，其中主机通过 JTAG 与设备通信。

上面，我们有一个 STM32F746ZG 开发板，它搭载了一个 ARM Cortex-M7 处理器，鉴于其在低功耗范围内具有强大的性能，它是边缘 AI 的理想选择。我们使用其 USB-JTAG 端口将其连接到我们的桌面机器。在桌面上，我们运行 OpenOCD 以打开与设备的 JTAG 连接；反过来，OpenOCD 允许 µTVM 使用设备无关的 TCP 套接字来控制 M7 处理器。有了这个设置，我们可以使用如下所示的 TVM 代码运行 CIFAR-10 分类器（完整脚本此处）

OPENOCD_SERVER_ADDR = '127.0.0.1'
OPENOCD_SERVER_PORT = 6666
TARGET = tvm.target.create('c -device=micro_dev')
DEV_CONFIG = stm32f746xx.default_config(OPENOCD_SERVER_ADDR, OPENOCD_SERVER_PORT)

module, params = get_cifar10_cnn()
with micro.Session(device_config) as sess:
	graph, c_module, params = relay.build(module['main'], target=TARGET, params=params)
  micro_mod = micro.create_micro_mod(c_module, DEV_CONFIG)
  graph_mod = graph_runtime.create(graph, micro_mod, ctx=tvm.micro_dev(0))
  graph_mod.run(data=data_np)
  prediction = CIFAR10_CLASSES[np.argmax(graph_mod.get_output(0).asnumpy())]
  print(f'prediction was {prediction}')

以下是 MicroTVM 的性能结果，与 CMSIS-NN 5.7.0 版本（提交 a65b7c9a），一个手动优化的 ML 内核库进行比较。

/images/microtvm/post-2020-05-28/cifar10-int-8-cnn.png

正如我们所见，开箱即用的性能并不是很好，但这就是 AutoTVM 发挥作用的地方。我们可以为我们的设备编写一个调度模板，进行一轮自动调优，然后获得明显更好的结果。要插入我们自动调优的结果，我们只需要替换这一行

graph, c_module, params = relay.build(module['main'], target=TARGET, params=params)

为以下几行

with TARGET, autotvm.apply_history_best(TUNING_RESULTS_FILE):
  graph, c_module, params = relay.build(module['main'], target=TARGET, params=params)

现在我们的结果看起来像这样

/images/microtvm/post-2020-05-28/autotuned-cifar10-int-8-cnn.png

我们的性能提高了约 2 倍，现在更接近 CMSIS-NN。虽然 MicroTVM CIFAR10 的实现与类似的 TFLite/CMSIS-NN 模型相比具有竞争力，但这项工作才刚刚开始利用 TVM 的优化功能。通过加速其他运算符（如密集/全连接层）并利用 TVM 的模型特定量化和运算符融合功能，还有进一步优化的空间。TVM 与 µTVM 使您能够充分发挥它们的优势。那么它是如何工作的呢？幕后发生了什么？现在让我们深入了解。

设计

/images/microtvm/post-2020-05-28/memory-layout.png
µTVM 设备在 RAM 中的内存布局

µTVM 旨在通过最小化必须满足的要求集来支持设备的最低公分母。特别是，用户只需提供

适用于其设备的 C 交叉编译器工具链
一种用于读取/写入设备内存并在设备上执行代码的方法
包含设备内存布局和一般架构特性的规范
一个准备设备以执行函数的代码片段

大多数裸机设备都支持 C 和 JTAG（一种调试协议），因此 (1) 和 (2) 通常是免费提供的！此外，(3) 和 (4) 通常是非常小的要求。以下是 STM32F746 系列开发板的 (3) 和 (4) 的示例。

device_config = {
    'device_id': 'arm.stm32f746xx',        # unique identifier for the device
    'toolchain_prefix': 'arm-none-eabi-',  # prefix of each binary in the cross-compilation toolchain (e.g., arm-none-eabi-gcc)
    'base_addr': 0x20000000,               # first address of RAM
    'section_sizes': {                     # dictionary of desired section sizes in bytes
         'text': 18000,
         'rodata': 100,
         'data': 100,
         ...
    },
    'word_size': 4,                        # device word size
    'thumb_mode': True,                    # whether to use ARM's thumb ISA
    'comms_method': 'openocd',             # method of communication with the device
    'server_addr': '127.0.0.1',            # OpenOCD server address (if 'comms_method' is 'openocd')
    'server_port': 6666,                   # OpenOCD server port (if 'comms_method' is 'openocd')
}

.syntax unified
.cpu cortex-m7
.fpu softvfp
.thumb

.section .text.UTVMInit
.type UTVMInit, %function
UTVMInit:
  /* enable fpu */
  ldr r0, =0xE000ED88
  ldr r1, [r0]
  ldr r2, =0xF00000
  orr r1, r2
  str r1, [r0]
  dsb
  isb
  /* set stack pointer */
  ldr sp, =_utvm_stack_pointer_init
  bl UTVMMain
.size UTVMInit, .-UTVMInit

µTVM 基础设施和设备运行时构建的目的仅在于利用这些要求，并且我们正在努力通过支持常见的开源运行时平台（如 mBED OS）来处理编译和链接过程，从而减少这些要求。

设备会话

鉴于微控制器交互的网络性质，我们通过引入 MicroSession 的概念，稍微偏离了标准的 TVM 代码。

µTVM 中的每个功能都依赖于与目标设备建立开放会话。如果您熟悉 TVM，您可能已经注意到我们第一个代码片段中有一行代码偏离了规范——即这一行

...
with micro.Session(device_config) as sess:
	...

此 with 代码块内的每一行都可以调用 µTVM 中的函数，上下文是由 device_config 指定的设备。这一行在幕后做了很多事情，所以让我们来拆解它。

首先，它使用您指定的任何通信方法（通常是 OpenOCD）初始化与设备的连接。然后，使用您指定的任何交叉编译器对 µTVM 设备运行时进行交叉编译。最后，主机分配编译后的二进制文件的空间，并使用打开的连接将二进制文件加载到设备上。

现在运行时已位于设备上，我们自然希望通过它运行一些函数。

模块加载

模块是 TVM 中的核心抽象之一。模块存储一组用于特定设备/运行时目标的相关函数。鉴于微控制器通常没有操作系统，µTVM 需要做很多额外的工作来维护这种高级抽象。为了了解发生了什么，我们将跟踪创建和加载 µTVM 兼容模块的过程。

假设我们有一个与设备打开的 micro.Session 和一个实现 2D 卷积的 TVM 调度。如果我们想将其加载到我们的微控制器上，我们需要它发出 C 代码。为此，我们只需要在 tvm.build 或 relay.build 中设置 target。示例

graph, c_module, params = relay.build(module['main'], target='c -device=micro_dev', params=params)

通过像这样设置目标，构建过程将通过我们的 C 代码生成后端运行。但是，生成的 C 模块仍然驻留在主机上。为了将其加载到设备上，我们通过 µTVM 基础设施中的核心函数之一运行它：create_micro_mod。示例

micro_mod = micro.create_micro_mod(c_module, DEV_CONFIG)

上面的代码行交叉编译模块内的 C 源代码，为生成的二进制文件分配空间（以便它可以与设备内存中的运行时共存），然后将二进制文件的每个部分发送到设备上分配的插槽。一旦模块二进制文件安全地位于设备内存中，就会修补二进制文件中的函数指针，以使模块能够访问设备运行时中的辅助函数（例如，用于分配暂存区）。

现在，我们的内核已加载到设备上，我们可以像这样获取卷积函数的远程句柄

micro_func = micro_mod['conv2d']

张量加载

如果我们想调用一个运算符，我们首先需要一些张量作为参数

data_np, kernel_np = get_conv_inputs()
ctx = tvm.micro_dev(0)
data = tvm.nd.array(data_np, ctx=ctx)
kernel = tvm.nd.array(kernel_np, ctx=ctx)

根据其数据类型（例如，int8、float32 等）和形状，计算每个张量的大小（以字节为单位），并且主机在设备的堆上分配一个内存区域。然后将张量的数据加载到分配的区域中。

函数调用

运算符执行可能是该系统中最棘手的部分。为了简化其介绍，我们将首先介绍严格执行（运算符在被调用后立即执行），然后是惰性执行（运算符仅在其结果需要时才执行）——后者是系统实际的工作方式。

严格执行

调用函数时，输入和输出张量都作为参数传递，这被称为目标传递风格

conv2D(data, kernel, output)

鉴于这些张量已在设备上分配，我们只需要将元数据发送到设备（设备地址、形状和数据类型），以便它知道要使用哪些驻留张量。函数调用的运行时表示包括此元数据，以及被调用函数的地址（如下所示）。在构建此表示之前，需要将元数据序列化到设备上专门为此目的存在的参数部分中。

/*
 * task struct for uTVM
 */
typedef struct {
  /* pointer to function to call for this task */
  int32_t (*func)(void*, void*, int32_t);
  /* array of argument tensors */
  TVMValue* arg_values;
  /* array of datatype codes for each argument */
  int* arg_type_codes;
  /* number of arguments */
  int32_t num_args;
} UTVMTask;

在严格设置中，存在一个全局 UTVMTask 实例，我们从主机端写入该实例。一旦我们写入任务，运行时就拥有执行函数所需的一切，我们就可以在运行时的入口点开始执行。运行时将执行一些轻量级初始化，运行我们的运算符，然后将控制权返回给主机。

惰性执行

在实践中，在用户请求后立即执行运算符会变得非常昂贵，因为通信开销开始占据主导地位。我们可以通过延迟评估直到用户想要调用结果来提高系统的吞吐量。

从实现的角度来看，我们现在需要在主机端累积函数调用元数据，而不是急于序列化参数元数据和 UTVMTask 数据，然后再将其刷新到设备。设备运行时还需要进行一些更改：（1）我们现在必须有一个全局 UTVMTask 数组，并且（2）我们需要循环并按顺序执行每个任务。

AutoTVM 与 MicroTVM

到目前为止，我们描述的运行时对于模型部署似乎不是很有用，因为它非常依赖主机。这是有意的，实际上，运行时是为不同的目标而设计的：AutoTVM 支持。

一般来说，AutoTVM 提出候选内核，在目标后端使用随机输入运行它们，然后使用计时结果来改进其搜索过程。鉴于 AutoTVM 只关心单个运算符的执行，我们将运行时设计为面向运算符的，而不是面向模型的。但在 µTVM 的情况下，与设备的通信通常会占据执行时间的大部分。惰性执行允许我们多次运行相同的运算符而无需将控制权返回给主机，因此通信成本在每次运行中均摊，并且我们可以更好地了解性能概况。

由于 AutoTVM 需要对大量候选内核进行快速迭代，因此 µTVM 基础设施目前仅使用 RAM。但是，对于自托管运行时，我们肯定需要同时使用闪存和 RAM。

托管图运行时

虽然托管运行时是为 AutoTVM 设计的，但我们仍然可以运行完整的模型（只要它们没有任何控制流）。只需使用 TVM 的图运行时，但在 µTVM 上下文中，就可以免费获得此功能。实际上，图运行时对主机的唯一依赖是张量分配和运算符调度（这只是依赖关系图的拓扑排序）。

评估

有了这个基础设施，我们试图回答以下问题

µTVM 真的与设备无关吗？
使用 µTVM 尝试优化需要多少工作量？

为了评估 (1)，我们在两个目标上进行了实验

一个 Arm STM32F746NG 开发板，配备 Cortex-M7 处理器
µTVM 主机模拟设备，它在主机上创建一个内存区域，并像裸机设备一样进行接口。

为了评估 (2)，我们探索了 Arm 开发板的优化，这些优化可以最大限度地提高性价比。

作为比较，我们从 Arm 的本教程中提取了一个量化的 CIFAR-10 CNN。在本教程中，CMSIS-NN（Arm 专家提供的高度优化的内核库）被用作运算符库，这使得这个 CNN 成为完美的评估目标，因为我们现在可以直接比较 µTVM 和 CMSIS-NN 在 Arm 开发板上的结果。

/images/microtvm/post-2020-05-28/cifar10-graphical.png
CIFAR-10 CNN 图示

方法

在我们的实验中，我们使用了来自 HEAD 的 TVM（提交 9fa8341）、CMSIS-NN 5.7.0 版本（提交 a65b7c9a）、STM32CubeF7 1.16.0 版本以及来自 Arm 的 GNU Tools for Arm Embedded Processors 9-2019-q4-major 9.2.1 工具链的 GCC（修订版 277599）。我们实验中使用的主机运行 Ubuntu Linux 18.04.4 LTS，并配备了 AMD Ryzen Threadripper 2990WX 32 核处理器，内存为 62GB RAM。此博客文章的所有评估脚本都包含在此仓库中。

Arm 特定优化

使用 CMSIS-NN，第一个卷积映射到他们的 RGB 卷积实现（专门用于输入层），而后两个卷积映射到他们的 “快速”卷积实现。我们认为在早期的通用优化之后，我们的 RGB 卷积性能已经足够接近，但对我们的快速卷积结果仍不满意。幸运的是，Arm 发布了一篇论文，描述了 CMSIS-NN 中使用的优化，我们发现他们通过 SIMD 内在函数获得了巨大的加速。在论文中，他们提出了一个使用 SIMD 内在函数的矩阵乘法微内核（如下图所示）。虽然我们可以在 TVM 的代码生成工具中添加对内在函数的一流支持——这可能是长期的最佳选择——但 TVM 提供了张量化作为支持 SIMD 的“快速而脏”的解决方案。

/images/microtvm/post-2020-05-28/simd-diagram.png
CMSIS-NN 论文中的图示，展示了一个 2x2 矩阵乘法微内核

张量化的工作原理是定义一个可以插入到 TVM 运算符最内层循环中的微内核。使用这种机制，为 Arm 开发板添加 SIMD 支持就像在 C 中定义一个微内核一样简单（此处），它镜像了他们论文中的实现。我们定义了一个使用此微内核的调度（此处），对其进行了自动调优，然后获得了“µTVM SIMD 调优”的结果。

虽然我们能够使用 SIMD 微内核进行直接卷积，但 CMSIS-NN 使用他们所谓的“部分 im2col”作为他们的实现策略，这在性能和内存使用之间提供了权衡。部分 im2col 不是一次性显现整个 im2col 矩阵，而是一次只生成几列。然后，对于每个批次，他们可以将矩阵发送到他们的 SIMD 矩阵乘法函数。

我们的假设是，在其他优化中，我们可以通过自动调优找到最佳批次大小。在实践中，我们发现部分 im2col 比我们的直接卷积实现慢得多，因此我们没有将其包含在其余结果中。

当然，我们还可以从 CMSIS-NN 中提取其他优化措施，以进一步缩小差距

将 int8 权重批量扩展为 int16，以减少 SIMD 的重复扩展
将卷积分割成 3x3 块以减少填充检查

但我们在这篇博客文章中的目标是展示可以使用 µTVM 完成的工作的概貌。即便如此，这也不是一场竞赛，因为 CMSIS-NN（以及任何其他手动优化的库）可以使用自带代码生成框架直接插入 TVM。

端到端

CIFAR-10

在探索了卷积的优化之后，我们开始测量它们对端到端性能的影响。对于 Arm 开发板，我们收集了未调优的结果、未使用任何 SIMD 调优的结果、使用 SIMD 调优的结果以及使用 CMSIS-NN 的结果。对于模拟主机设备，我们仅收集了未调优的结果和通用调优的结果。

https://github.com/areusch/microtvm-blogpost-eval

/images/microtvm/post-2020-05-28/autotuned-cifar10-int-8-cnn.png
int8 量化 CIFAR-10 CNN 在 Arm STM32F746NG 上的比较（从上方重新发布）

/images/microtvm/post-2020-05-28/autotuned-cifar10-int-8-cnn-x86.png
int8 量化 CIFAR-10 CNN 在 µTVM 的模拟主机设备上的比较

在 Arm STM32 系列开发板上，与最初的未调优运算符相比，我们能够将性能提高约 2 倍，并且我们获得的结果更接近 CMSIS-NN。此外，我们还能够在主机模拟设备上显着提高性能。虽然 x86 数字意义不大，但它们表明我们可以使用相同的基础设施 (µTVM) 来优化截然不同的架构上的性能。

请继续关注未来更多的端到端基准测试，因为我们将更广泛地扩展这种方法。

自托管运行时：最后的边疆

/images/microtvm/self-hosted-runtime.png

设想的 µTVM 优化和部署流程

虽然如上所述，使用当前的运行时已经可以获得端到端基准测试结果，但在独立模式下部署这些模型目前仍在我们的路线图上。差距在于，面向 AutoTVM 的运行时目前依赖主机来分配张量和调度函数执行。但是，为了在边缘端发挥作用，我们需要一个通过 µTVM 的流程，该流程生成一个单一的二进制文件，以便在裸机设备上运行。然后，用户可以通过在其边缘应用程序中包含此二进制文件，轻松地将快速 ML 集成到他们的应用程序中。此流程的每个阶段都已经到位，现在只是将它们粘合在一起的问题，因此请期待我们很快在此方面发布更新。

结论

用于单内核优化的 MicroTVM 今天就可以使用，并且是该用例的最佳选择。随着我们现在构建自托管部署支持，我们希望您和我们一样兴奋，使 µTVM 也成为模型部署的最佳选择。然而，这不仅仅是一项观赏性运动——请记住：这一切都是开源的！µTVM 仍处于早期阶段，因此每个人都可以对其发展轨迹产生重大影响。如果您有兴趣与我们一起构建，请查看 TVM 贡献者指南，或者直接进入 TVM 论坛首先讨论想法。

致谢

这项工作的完成离不开以下人员的贡献

Tianqi Chen，感谢他在设计方面的指导，以及作为一位出色的导师。
Pratyush Patel，感谢他在 MicroTVM 早期原型上的合作。
OctoML，感谢他们为我提供的实习机会，使我能够全力投入到这个项目中。
Thierry Moreau，感谢他在我在 OctoML 工作期间对我的指导。
Luis Vega，感谢他教我与微控制器交互的基础知识。
Ramana Radhakrishnan，感谢他提供我们实验中使用的 Arm 硬件，并提供有关其使用的指导。