TinyML-TVM是如何驯服Tiny的（上）

TinyML-TVM是如何驯服Tiny的（上）

低成本、人工智能驱动的消费类设备的激增，导致了ML研究人员和从业者对“裸智能”（低功耗，通常没有操作系统）设备的广泛兴趣。虽然专家已经可以在一些裸机设备上运行某些模型，但是为不同设备集优化模型是一个挑战，通常需要手动优化特定于设备的库。对于那些没有Linux支持的平台，没有可伸缩的模型部署解决方案。正因为如此，为了瞄准新设备，开发人员必须实现一次性定制软件堆栈，以管理系统资源和调度模型执行。             

机器学习软件的手动优化并不是裸机领域所独有的。事实上，这一直是与其他硬件后端（例如，gpu和fpga）一起工作的开发人员的共同主题。TVM已经被证明能够抵御新硬件目标的冲击，但直到现在，它还无法与微控制器的独特配置相抗衡。为了解决这个领域的问题，扩展了TVM，使其具有一个微控制器后端，称为µTVM（脚注：发音为“MicroTVM”）。µTVM有助于在裸机设备上执行tensor程序，并通过TVM的内置tensor程序优化器AutoTVM自动优化这些程序。下图显示了µTVM+AutoTVM基础设施的鸟瞰图：

 

让看看它的行动             

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

 

 A standard µTVM setup, where the host communicates with the device via JTAG.

上面，有一个STM32F746ZG板，里面有一个ARM Cortex-M7处理器，考虑到它在低功耗封装中的强大性能，这是边缘人工智能的理想部件。使用它的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（commit a65b7c9a）相比，后者是一个手工优化的ML内核库。

 

开箱即用的性能不是很好，但这正是AutoTVM的救命稻草。可以为设备编写一个调度模板，进行一轮自动调整，然后获得显著更好的结果。要插入自动调谐结果，只需要替换这一行：

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

with these lines:

with TARGET, autotvm.apply_history_best(TUNING_RESULTS_FILE):

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

And our results now look like this:

 

性能提高了约2倍，现在离CMSIS-NN更近了。尽管MicroTVM CIFAR10的实现与类似的TFLite/CMSIS-NN模型相比具有竞争力，但这项工作刚刚开始利用TVM的优化特性。通过加速其他运营商（如密集/全连接dense/fully-connected）和利用TVM的模型特定量化和运算符融合功能，还有进一步优化的空间。带有µTVM的TVM能够发挥最佳性能。如何工作的呢？幕后是怎么回事？现在就开始吧。

Design

 

The µTVM Device Memory Layout in 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代码。             

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

...

with micro.Session(device_config) as sess:

    ...

此with块内的每一行都可以调用µTVM中的函数，上下文是device_config指定的设备。这条线在hood下面做了很多事情，让把它拆开。              

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

由于 runtime现在位于设备上，自然需要一些函数来运行它。

Module Loading

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

假设有一个微型会议打开设备和实现二维卷积的TVM调度。如果想把它加载到微控制器上，需要它发出C代码。要做到这一点，只需要设定目标tvm.build or relay.build. Example:

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访问help函数（例如，分配草稿行）。             

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

micro_func = micro_mod['conv2d']

Tensor Loading

If we want to call an operator, we first need some tensors as arguments:

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)             

考虑到这些张量已经在设备上分配，只需要向设备发送元数据(device address, shape, and data type)（设备地址、形状和数据类型），这样设备就知道要使用哪个驻留张量。下面显示的是一个名为“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将执行一些轻量级初始化，运行算子，然后将控制权返回给主机。