TVM如何训练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基础架构和设备Runtime仅用于满足这些要求，正在努力通过支持通用的开源Runtime平台（例如mBED OS）来处理编译和链接过程，以降低这些要求。

设备会话

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

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

...

with micro.Session(device_config) as sess:

        ...

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

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

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

模块加载

TVM中的核心抽象之一是模块的抽象。模块存储用于特定设备/Runtime目标的一组相关功能。鉴于微控制器通常没有操作系统，因此µ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源代码，为所得的二进制文件分配空间（以便可以与Runtime在设备内存中共存），然后将二进制文件的每个部分发送到其在设备上分配的插槽中。一旦模块二进制文件贴紧在设备内存中，便会修补二进制文件中的功能指针，以使模块可以在设备Runtime访问辅助功能（例如，用于分配暂存器）。

现在，在将内核加载到设备上之后，可以获取卷积函数的远程句柄，如下所示：

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)

鉴于这些张量已在设备上分配，只需要将元数据发送到设备（设备地址，形状和数据类型），知道要使用哪个驻留张量。函数调用的Runtime表示形式包括此元数据以及被调用函数的地址（如下所示）。在构造此表示形式之前，需要将元数据序列化到为此目的明确存在的设备上的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实例，从主机端将其写入其中。一旦写入任务，Runtime就具有执行该功能所需的一切，并且可以在Runtime的入口点开始执行。Runtime将执行一些轻量级的初始化，运行算子，然后将控制权返回给主机。

执行

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

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

带MicroTVM的AutoTVM

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

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

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

托管图Runtime

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

评估

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

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

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

• 一个手臂STM32F746NG开发板，采用了的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核处理器。

特定于手臂的优化

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

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

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

尽管能够使用SIMD微内核进行直接卷积，但是CMSIS-NN使用他们所谓的“部分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）来在极为不同的体系结构上优化性能。

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

自托管Runtime：最终领域

 

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

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

结论

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