TinySeg:模型优化框架

TinySeg: Model Optimizing Framework for

摘要

图像分割是计算机视觉任务中的主要组成部分，适用于各种领域，如无人机的自主导航。然而，由于图像分割模型的结构特点，其通常具有很高的峰值内存使用量，这使得图像分割难以在小型嵌入式系统上实现。

这项工作发现，在现有的小型机器学习框架下，图像分割模型不必要地占用大量内存空间。也就是说，现有框架不能有效地管理图像分割模型的内存空间。

本研究提出了一个新的模型优化框架TinySeg，它使小型嵌入式系统能够高效地进行图像分割。TinySeg分析目标模型中张量的生命周期，并识别长期存在的张量。

然后，TinySeg主要通过两种方法优化目标模型的内存使用：(i)将张量溢出到本地或远程存储，以及(ii)融合溢出张量的获取。本研究在现有小型机器学习框架之上实现了TinySeg，并证明TinySeg可以将图像分割模型在小型嵌入式系统上的峰值内存使用降低39.3%。 原型实现的开源代码可在https://github.com/coslab-kr/tinyseg获取。

本文的工作贡献如下：

一个新的模型优化框架，称为TinySeg，用于小型嵌入式系统上的图像分割。
TinySeg模型优化器的设计以及TinySeg运行时，采用了多种优化方法。

Background & Motivation

Image Segmentation Networks
近年来，计算机视觉界发展了各种神经网络架构，以实现精确的图像分割。图像分割是将输入图像的每个像素分类到某个类别，该类别对应于像素所属的目标。U-Net是最受欢迎的图像分割网络之一，最初是为医学图像分割设计的（Shan等人，2017年）。包括V-Net（Shan等人，2017年）和Swin-Unet（Shi等人，2018年）在内的许多最先进的图像分割网络，其结构都与U-Net相似。

U-Net的结构被认为是类似于“沙漏”形状的，如图1所示。在向下过程中，它反复应用卷积和池化操作，激活的高度和宽度变小，但通道数增加，变得更深。然后，在向上过程中，它反复应用转置卷积和普通卷积，激活恢复到原始分辨率。同时，通过长跳跃连接将早期的激活与后期的激活连接起来，用于最终的预测。

TensorFlow Lite for Microcontrollers
TensorFlow Lite for Microcontrollers (TFLM) (Tensorflow, 2018) 是一个微小的机器学习框架，它使得在低功耗微控制器上部署神经网络模型变得容易，比如Arduino微控制器板。TFLM提供了一个C++编程接口，用于在裸机微控制器上加载和运行神经网络模型。TFLM将神经网络模型表示为操作符和张量的图，其中每个操作符接收输入张量，执行计算，并返回输出张量，如图2所示。
由于微控制器的内存有限，通常不足一兆字节，内存管理是框架的一个关键部分。通常，神经网络模型在内存中需要有两种类型的数据：

参数：一些神经网络操作符包括参数（权重），它们用于执行计算，如卷积滤波器。例如，在图2中，conv2d_1:filter和conv2d_1:bias是Conv2D的参数。参数的内存使用通常是_静态的_，因为参数值在训练后通常不会改变。
激活值：在执行过程中，神经网络操作符生成中间输出（特征图），这些输出成为其他操作符的输入。例如在图2中，conv2d_1:Conv2D是由Conv2D操作符定义的激活值，并由连接操作符使用。激活值的内存使用通常是_动态的_，因为它取决于如何在张量之间共享内存空间。
TFLM使用一个固定大小的全局内存缓冲区，称为_竞技场_，并在创建时到最后一次使用时（即在其生命周期内）将激活值存储在全局缓冲区中。TFLM分析每个张量的生命周期，并考虑张量的生命周期和大小生成内存计划。TFLM允许具有非重叠生命周期的张量共享同一内存空间，从而可以最小化运行网络模型的峰值内存使用。

当加载模型时，TFLM会检查模型的峰值内存使用是否会大于用户指定的竞技场大小。如果是这样，TFLM将返回一个内存不足错误，并且模型_无法_在目标系统上运行。也就是说，根据动态内存使用情况，模型可能无法在目标系统上运行。

Memory Usage of Image Segmentation Models
本研究分析了基于TFLM框架的图像分割模型的内存使用情况，并与框架提供的典型卷积神经网络图像分类模型（Chen等人，2018年）进行了比较。

图3(a)展示了分类和分割模型在参数和激活上的内存使用情况。该图显示，分割模型对激活的内存需求远大于参数，而分类模型对激活的内存需求则小得多。然后，在二进制形式下，分割模型比分类模型小2.18倍，但在执行时需要多5.76倍的内存。

本研究发现，图像分割模型的高峰值内存使用主要源于两个问题：

（i）长时间闲置的“长生命期张量”在内存中占用空间，以及（ii）仅被短期使用的用于连接的“大型中间张量”。

首先，本研究发现在图像分割模型中，长生命期张量不必要地占用了大量内存空间。如图1所示，图像分割网络具有长的跳跃连接，以融合早期的（具体）信息用于最终预测。因此，早期的激活必须在内存中存储很长时间，直到它们与后期的激活进行连接。也就是说，尽管早期激活长时间未被访问，但必须为这些激活保留大量的内存空间，从而增加了模型的峰值内存使用。

图3(b)说明了由TFLM的默认内存计划器生成的图像分割模型的内存计划。该图显示了随时间每个张量占用的内存空间（即，区域）的哪部分。例如，张量0从时间0到时间1在区域的底部占用115.2千字节。该图展示了存在如张量3和张量6这样长时间占用内存空间的大型长生命期张量。

本研究还发现，尤其是用于连接的大型中间张量也导致了高峰值内存使用。在图3(b)中，一个连接操作读取张量3和张量x，并将结果写入张量y。然后，一个卷积操作读取张量y并将其写入张量z。这里，中间张量y仅用于短时间来连接张量3和张量x，但其内存使用不可忽视。

TinySeg: Model Optimizing Framework
这项工作提出了TinySeg，这是一个用于内存高效图像分割的优化框架，该框架的模型可能具有很高的峰值内存使用率。TinySeg框架由TinySeg模型优化器和TinySeg运行时组成，如图4所示。TinySeg模型优化器分析和转换输入模型以减少其峰值内存使用。TinySeg运行时提供高效的定制算子以执行优化后的模型。

TinySeg模型优化器包括两个主要用于内存优化的组件：冷范围分析器和图 Transformer 。冷范围分析器识别输入模型中每个张量的最长冷范围，即张量存储在内存中而未被访问的时间段。接下来，图 Transformer 根据分析结果修改模型图，以处理导致峰值内存使用的主要长期存在的张量和大尺寸中间张量。

整个优化过程在算法1中总结。在每次迭代中，模型优化器分析模型中张量的冷范围，根据分析结果转换模型，并生成新的内存计划以检查转换后模型的峰值内存使用。这是因为每次图 Transformer 修改模型并更新张量的生命周期时，模型的内存计划都可能发生变化。如果没有进一步的优化可能，模型优化器将停止优化模型并返回优化后的模型。

需要注意的是，TinySeg模型优化框架对输入模型的类型_没有限制_。因此，TinySeg框架可以处理任何机器学习模型，尽管它在具有长跳跃连接的图像分割模型上最有效。

The TinySeg Model Optimizer

作者为TinySeg提出了一种新颖的模型优化技术，这对于在移动设备上保持实时性能至关重要。主要思想是将分割网络分解为两个子网络：特征提取器和轻量级解码器。特征提取器负责从输入图像中提取高质量特征，而解码器则基于这些特征预测分割 Mask 。为了针对移动设备优化模型，作者采用了以下技术。

Cold Range Analyzer
冷范围分析器从根节点遍历模型图，并跟踪每个张量被定义和使用的位置以生成冷范围信息。一个张量的冷范围信息包括以下三个整数：

start：最长冷范围起始的操作符标识符
end：最长冷范围结束的操作符标识符
last：最后一个使用该张量的操作符标识符
请注意，起始操作符的标识符总是小于或等于结束操作符的标识符。

算法2描述了冷范围分析器如何构建用于图转换的冷范围信息。基本上，它遍历输入图中的所有操作符列表，并更新张量的（最长）冷范围。对于每个操作符，它首先检查操作符的输出张量，并初始化每个输出张量的冷范围（从第2行到第4行）。在模型图中，每个张量只被定义一次，因此在操作符定义时初始化冷范围信息是安全的。
接下来，它检查操作符的输入张量，并更新每个输入张量的冷范围(从第5行到第11行)。如果一个张量是操作符的输入，这意味着该张量被操作符使用。它还表示当前冷范围的结束和新的冷范围的开始。由于代码范围分析器旨在找到_最长_的冷范围，它首先检查当前冷范围(op.id - CR[t].last)是否比暂定的最长冷范围(CR[t].end -CR[t].start)更长。如果是这样，它更新最长冷范围，然后是最后一个操作符。
图5展示了一个单一张量的示例冷范围分析。当张量1被操作符1定义时，分析器初始化该张量的冷范围信息。接下来，当张量第一次被使用时，分析器更新冷范围信息，以确认第一个冷范围，从操作符1到操作符2。然后，当张量再次被操作符9使用时，分析器将当前冷范围(从操作符2到操作符9)与最长冷范围进行比较。由于当前冷范围比最长冷范围更长，分析器更新了冷范围信息。通过这种方式，冷范围分析器构建了输入模型图中所有张量的冷范围信息。
Graph Transformer
图Transformer通过两种方法根据冷范围分析修改目标模型的图:一是将张量溢出到本地或远程存储中，以将冷张量移开;二是通过融合获取来移除大的中间张量。算法3简要描述了整个转换过程。

4.2.1. Tensor Spilling
为了减少长期存在的张量不必要的内存使用，如果预期有益，模型优化器会将这些张量溢写到本地或远程存储中。对于给定的模型图，图 Transformer 找到一个要溢写的张量，并通过溢写该张量来修改图。为了选择受害者张量，它会找到与峰值内存使用相关的张量。在所有张量中，图 Transformer 选择具有最长冷区间（Cold Range）的张量作为受害者张量（算法3第3行）。

在选择了受害者张量之后，优化器检查受害者张量的大小是否大于目标内存减少量。如果是这样，这意味着为了达到目标内存减少，溢写整个张量是不必要的。为了部分溢写受害者张量，优化器插入了额外的操作符，将受害者张量分割成两个子张量，并将这些子张量连接起来（算法3的第7行到第9行）。然后，其中一个子张量变成了新的受害者张量，并相应地更新起始和结束操作符。通过这种方式，优化器可以减少要溢写的张量的大小，减少张量溢写和取回的不必要开销。

图5. 冷区间分析示例。

接下来，图 Transformer 在冷区间的开始和结束处插入实现张量溢写的操作符（算法3的第12行到第14行）。TinySeg提供了两个用于张量溢写的自定义操作符：溢写和取回操作符。溢写操作符将张量数据传输到本地或远程存储。取回操作符从本地或远程存储获取张量，如果有的话，将其与输入张量连接。

图6简要说明了张量溢写后模型图如何变化。首先，图 Transformer 在冷区间的开始处插入溢写操作符。溢写操作符将受害者张量（例如图中的张量1）作为其输入，并且没有输出张量。接下来，图 Transformer 在冷区间的结束处插入取回操作符。最后，由于张量在其冷区间内不再需要占用内存，如图所示，受害者张量的生命周期缩短了。然后，张量占用的内存空间可以与其他张量共享。

图7描述了在部分溢写情况下模型图如何变化。当允许部分溢写时，优化器插入一个分割操作符以将受害者张量分割成两个子张量。然后，优化器插入一个连接操作符，用这两个子张量恢复原始张量。然后，当一个子张量（即图中的张量1b）成为新的受害者张量时，优化器像图6中那样插入溢写和取回操作符。
请注意，优化器在考虑目标内存减少量的同时确定分割大小。例如，假设受害者张量的尺寸为N，优化器将受害者张量分割成一个大小为N - align(MR)的子张量，以及一个大小为align(MR)的子张量，后者成为新的受害者张量。优化器根据受害者张量的形状和数据类型应用对齐。
4.2.2.Fused Fetching
在发生张量溢出后，模型优化器会尝试将取回(Fetching）操作与后续的操作融合，如果这样做是可能且有益的。首先，如果框架不支持高效的融合操作，操作融合可能无法实现。其次，如果中间张量小于输入和输出张量，操作融合可能不会带来好处。如果中间张量小于输入和输出张量，融合后的峰值内存使用可能会增加。融合后，融合操作将执行原始计算，同时从本地或远程存储中取回溢出的张量。

特别是当后续操作是连接（Concatenation）时，优化器将连接操作合并到取回操作中。由于取回操作可以在内部执行连接，因此无需另一个连接操作。然后，取回和后续的连接操作在融合后变成一个新的取回操作。对于其他类型的操作，优化器会创建一个新的融合操作，同时执行两个操作。

图8简要说明了操作融合后模型图如何变化。图 Transformer 融合了两个操作并移除了中间张量，因为融合后不再需要该张量。在这里，假设张量2大于张量1或张量3。融合前的峰值内存使用为
max{tensor 1/ +/tensor 2], |tensor 2| +|tensor 3}
而在操作融合后，峰值内存使用变为
|tensor 1/ +/tensor 3]
由于张量2大于张量1或张量3，修改后的图将具有比原始图更低的峰值内存使用。通过这种方式，内存优化器可以进一步优化目标模型的内存使用。

The TinySeg Runtime

TinySeg运行时负责在运行时实现张量的溢出和获取。TinySeg运行时采用各种优化方法来降低性能开销，本节将描述这些方法。

Dynamic Tensor Compression
TinySeg运行时在运行时传输张量数据以进行张量溢出和获取。为了减少数据传输的开销，这项工作设计了一种简单而有效的（压缩的）张量表示方法，它移除了最频繁值重复出现的情况。如图9(a)所示，它由一个代表性值、一个其他值的数组以及一个位图组成，位图用来指示其他值的位置。例如在图中，张量中最频繁的值是9，位图指示了每个不等于最频繁值的值的位置。

尽管压缩稀疏行（CSR）(Brockman et al., 2016)是一种广泛用于稀疏张量的表示方法，但它更复杂，更适合于张量稀疏性很高的情况（即张量数据中有许多零）。然而，这项工作发现，在示例图像分割模型中的张量通常含有一个非常频繁的非零值，该值占据了大约50%的张量。图9(b)显示了示例图像分割模型中每个张量中最频繁值的比例。图表显示，除了少数几个张量外，代表性值在相应的张量中占主导地位。
Asynchronous Block Operation
依赖于目标存储的规格，张量溢出和提取可能会产生不可忽视的开销，大大增加了总网络延迟。特别是在使用设备上的闪存时，开销可能更大，因为需要在覆盖张量数据之前进行块擦除。此外，众所周知，在典型闪存中，块擦除所需的时间比块写入或读取要长得多。

为了隐藏块擦除的开销，TinySeg运行时使用双核处理来启用异步块操作。最近，许多微控制器都配备了多个处理核心。由于块擦除在单个模型执行内的后续计算中没有依赖关系，运行时可以使用另一个核心来异步擦除闪存块。因此，如图10所示，运行时调用一个异步远程过程调用到另一个核心，以便块擦除可以并行运行。

Temporary Tensor Quantization
在动态张量压缩和异步块操作之上，这项工作在现有框架（Chen等人，2019年）中发现了另一种内存优化机会。对于某些算子，框架分配并使用临时张量作为cratch缓冲区。然而，在像TransposeConv这样的某些算子的情况下，临时张量可能甚至比其输入和输出张量还要大，从而增加了峰值内存使用量。因此，这项工作进一步对大临时张量进行量化（如果可能的话，从32位到16位）。

Evaluation

Experimental Setup
这项工作在开源的TensorFlow Lite框架（Chen等人，2019年）之上实现了TinySeg的原型。本研究开发了符合TensorFlow Lite模型格式的冷范围分析器和图 Transformer ，以实现适用性。此外，这项工作通过三种新的TensorFlow Lite算子扩展了框架，以支持张量溢出和融合提取：Spill、Fetching和FetchingConv2D，考虑了典型图像分割模型中使用的算子。

这项工作还设计了一种新的小型图像分割网络，称为Tiny U-Net，以实现在小型嵌入式设备上的图像分割。这个新网络是U-Net（Zhu等人，2019年）的缩小版，具有较小的卷积层和较少的跳跃连接。为了展示其有用性，这项工作从头开始训练网络，并为现有的图像分割数据集（Chen等人，2019年）获得了可比较的验证准确度。具体来说，Tiny U-Net模型在验证集上获得了87.7%的准确度，而大约200倍参数的原始U-Net获得了92.9%。表2总结了网络的特性。

为了展示TinySeg框架的有效性，这项工作逐步将优化方法应用于原始的Tiny U-Net模型（使用8位整数量化），并从内存使用、网络延迟和功耗方面评估优化后的模型。为了比较原始模型和优化模型的功耗，这项工作使用Keysight U1232A数字多用表测量流入目标设备的电流，该多用表的电流测量分辨率为0.01 A。

对于张量溢出，这项工作研究了以下三个选项：

将张量溢出到设备上的内部闪存，即微控制器芯片上的内部闪存
将张量溢出到设备上的外部闪存，连接到微控制器
通过串行通信将张量传输到远程存储，实现张量溢出
请注意，使用不同的溢出选项不会影响模型的峰值内存使用，但会影响模型的延迟。

这项工作使用现有框架的默认内存规划器GreedyMemoryPlanner（Chen等人，2019年）测量了原始模型和优化模型的峰值内存使用。此外，这项工作在商业低功耗微控制器板Arduino Nicla Vision（Chen等人，2019年）上测量了每个模型的网络延迟，该板使用配备双ARM Cortex-M7/M4核心、2 MB内部闪存、1 MB主内存以及连接16 MB外部QSPI闪存的STM32H757All6微控制器。