TinyML-TVM如何驯服TinyML

TinyML-TVM如何驯服TinyML

低成本，以人工智能为动力的消费类设备的激增，导致机器学习研究人员和从业人员对“裸机”（低功耗，通常没有操作系统）设备产生了广泛的兴趣。尽管专家已经有可能在某些裸机设备上运行某些模型，但是为各种设备优化模型的挑战非常艰巨，通常需要手动优化设备特定的库。对于那些没有Linux支持的平台，不存在用于部署模型的可扩展解决方案。因此，为了定位新设备，开发人员必须实现一次性的定制软件堆栈，以管理系统资源和调度模型执行。

机器学习软件的手动优化不是裸机设备领域独有的。实际上，对于那些使用其它硬件后端（例如GPU和FPGA）的开发人员来说，这是一个共同的主题。TVM已被证明可以抵御新硬件目标的冲击，但直到现在，它仍无法解决微控制器的独特特性。为了解决这一领域的问题，扩展了TVM以提供称为µTVM（脚注：发音为“ MicroTVM”）的微控制器后端。µTVM有助于在裸机设备上执行主机驱动的张量调度，并通过TVM内置的张量调度优化器AutoTVM来自动优化这些调度。下图显示了µTVM + AutoTVM基础架构的鸟瞰图：

作用

在讨论什么是TVM / MicroTVM或如何工作之前，看一下实际使用示例。

标准µ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}')

与CMSIS-NN 5.7.0版（commit a65b7c9a）（一种手动优化的ML内核库）相比，以下是MicroTVM的性能结果。

开箱即用的性能并不好，但这就是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)

现在，结果如下所示：

性能提高了约2倍，现在，离CMSIS-NN更近了。尽管MicroTVM CIFAR10实施方案与类似的TFLite / CMSIS-NN模型方面具有竞争力，但这项工作才刚刚开始利用TVM的优化功能。通过加速其它调度（如密集/完全连接）并利用TVM特定于模型的量化和调度融合功能，还有进一步优化的空间。带有µTVM的TVM能够充分发挥作用。它是怎样工作的？幕后发生了什么事？

设计

µTVM设备在RAM中的存储器布局

µTVM旨在通过最大限度地减少必须满足的一组要求，支持设备的最小公分母。需要提供：

1. 设备的C交叉编译器工具链
2. 一种用于读取/写入设备内存并在设备上执行代码的方法
3. 包含设备的内存布局和一般体系结构特征的规范
4. 一个代码片段，为设备执行功能做准备

大多数裸机设备都支持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）来处理编译和链接过程，以降低这些要求。

设备会话

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

µTVM中的每个功能都依赖于与目标设备的开放会话。如果熟悉TVM，可能会注意到有一行代码与第一个代码段中的规范有所不同-即，这是一个代码：

...

with micro.Session(device_config) as sess:

...

该with块内的每一行都可以调用µTVM中的函数，上下文是所指定的设备device_config。这条线在做很多事情，所以拆开包装。

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

有了设备上的运行时，自然会希望一些功能运行。

模块加载

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

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

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)

鉴于这些张量已在设备上分配，只需要将元数据发送到设备（设备地址，形状和数据类型），要使用哪个驻留张量。函数调用的运行时表示形式包括，此元数据以及被调用函数的地址（如下所示）。在构造此表示形式之前，需要将元数据序列化到目标明确存在的设备上的arguments部分中。

/*

 * 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实例，从主机端将其写入其中。一旦写入任务，运行时就具有执行该功能所需的一切，并且可以在运行时的入口处开始执行。运行时将执行一些轻量级的初始化，运行算子，然后将控制权返回给主机。

lazy执行

在实践中，由于通信开销开始占主导地位，一旦用户要求执行算子就变得非常昂贵。可以通过延迟评估直到用户希望获得调用结果的方式来提高系统的吞吐量。

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

带MicroTVM的AutoTVM

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

通常，AutoTVM会提出候选内核，并使用随机输入在目标后端运行，然后使用调度结果来改善其搜索过程。鉴于AutoTVM只关心单个算子的执行，将运行时设计为面向算子，而不是面向模型。但是对于µTVM，与设备的通信通常会占据执行时间。lazy惰性执行可以多次运行同一算子，而无需将控制权交还给主机，因此，通信成本在每次运行时均摊销，可以更好地了解性能概况。

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

托管图运行时

尽管托管的运行时是为AutoTVM设计的，但仍然可以运行完整的模型（没有任何控制流）。仅通过使用TVM的图形运行时，即可免费使用此功能，但具有µTVM上下文。实际上，图运行时对主机的唯一依赖是张量分配和算子调度（这只是依赖图的一种拓扑类型）。

评估

有了这个基础架构，试图回答以下问题：

1. µTVM是否真的与设备无关？
2. 使用µTVM进行优化试验需要多少算力？

为了评估（1），在两个目标上进行了实验：

• 一个armSTM32F746NG开发板，采用了的Cortex-M7处理器
• µTVM主机模拟设备，可以在主机上创建一个内存竞技场，与之连接的主机就像裸机设备一样。

为了评估（2），探索Arm板的优化方案，这些方案可以最大程度地降低成本。

作为比较，从Arm中提取量化的CIFAR-10 CNN 。CMSIS-NN（Arm专家高度优化的内核库）用作算子库，使该CNN成为理想的评估目标，可以直接将µTVM的结果与Arm上的CMSIS-NN进行比较木板。

CIFAR-10 CNN图

方法

在实验中，使用HEAD的TVM（commit 9fa8341），CMSIS-NN的5.7.0版（commit a65b7c9a），STM32CubeF7的1.16.0版以及Arm的适用于Arm嵌入式处理器的GNU工具的GCC 9-2019-q4-major 9.2 .1工具链（修订版277599）。实验中，使用的主机运行Ubuntu Linux 18.04.4 LTS，运行带有62GB RAM的AMD Ryzen Threadripper 2990WX 32核处理器。

特定于arm的优化

使用CMSIS-NN，第一个卷积映射到其RGB卷积实现（专门用于输入层），而后两个卷积映射到其“快速”卷积实现。经过较早的泛型优化后，性能对于RGB卷积已经足够接近了，但是对快速卷积结果却不满意。幸运的是，Arm发布了一篇描述CMSIS-NN中使用的优化的论文，发现正从SIMD内在函数中获得巨大的提速。本文提出了一种使用SIMD内在函数的矩阵乘法微内核（下图）。虽然可以在TVM的代码生成工具中添加对内在函数的一流支持，这从长远来看可能是最好的做法，TVM 张量化是支持SIMD的“快捷方法”。

CMSIS-NN论文的图表显示了2x2矩阵乘法微内核

张量化通过定义可插入TVM算子最内层循环的微内核来工作。使用这种机制，添加对Arm板的SIMD支持，就像在C中定义一个微内核一样简单，该微内核反映了其论文中的实现。定义了一个调度，使用该微内核，对其进行自动调整，然后得到“ µTVM SIMD调整”结果。

尽管能够使用SIMD微内核进行直接卷积，但是CMSIS-NN使用称为“ partial im2col”的实现策略，这在性能和内存使用之间进行了权衡。代替一次显示整个im2col矩阵，部分im2col一次只生成几列。然后，对于每一批，可以将矩阵发送到其SIMD matmul函数。

假设是，除其它优化外，可以通过自动调整找到最佳的批量大小。在实践中，发现部分im2col比直接卷积实现要慢得多，不在其余结果中。

当然，还可以从CMSIS-NN中获得其它优化，以进一步缩小差距：

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

本文目标是展示µTVM可以完成的工作的大致范围。即使这样，这也不是竞争，因为CMSIS-NN（以及任何其它手动优化的库），可以使用Bring Your Own Codegen框架直接插入TVM 。

端到端

CIFAR-10

在探索卷积优化之后，着手测量其对端到端性能的影响。对于ARM板，收集了未调整的结果，这是调整的结果没有任何使用SIMD，这是调整的结果与SIMD和结果使用CMSIS-NN。对于模拟的主机设备，仅收集未调整的结果和通用的调整结果。

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

int8Arm STM32F746NG进行量化的CIFAR-10 CNN比较

int8µTVM的仿真主机设备上对量化的CIFAR-10 CNN进行比较

在Arm STM32系列板上，与最初的未调整算子相比，能够将性能提高约2倍，并且所获得的结果更接近CMSIS-NN。此外，能够显着提高主机仿真设备上的性能。尽管x86的数字意义不大，但可以使用相同的基础架构（µTVM）来在极为不同的体系结构上优化性能。

随着更广泛地扩展，在将来继续关注更多端到端基准测试。

自托管运行时：最终领域

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

如上所述，虽然当前运行时已经可以获取端到端基准测试结果，但目前仍在路线图上以独立能力部署这些模型。差距在于面向AutoTVM的运行时当前依赖于主机来分配张量并调度函数执行。为了在边缘是使用，需要通过μTVM，其产生一个管道单一待裸机设备上运行的二进制。然后，用户可以通过将此二进制文件包含在边缘应用调度中，轻松将快速ML集成到其应用调度中。该管道的每个阶段都已经到位，现在只需将它们粘合在一起即可，期待在此方面的最新进展。

结论

用于单内核优化的MicroTVM现已准备就绪，并且是该用例的选择。现在，当建立自托管的部署支持时，希望也和使µTVM成为模型部署的选择一样。但是，这不只是一场观看比赛-记住：这都是开源的！µTVM仍处于起步阶段，因此每个人对其轨迹都会产生很大的影响。