MCUFormer: Deploying Vision Transformers on Microcontrollers with Limited Memory
- 论文链接：https://ziweiwangthu.github.io/data/MCUFormer.pdf
- 源码链接：https://hub.yzuu.cf/liangyn22/MCUFormer
简介
用于现实应用的深度神经网络部署通常需要高性能计算设备，如GPU和TPU。由于这些设备的高昂价格和能耗，不可接受的部署费用严格限制了深度模型在各种任务中使用。用于现实应用的深度神经网络部署通常需要高性能计算设备，如GPU和TPU。由于这些设备的高昂价格和能耗，不可接收的部署费用严格限制了深度模型在各种任务中使用，尤其是在没有足够电池支持的场景中。在低能耗的廉价物联网设备上部署深度神经网络称为许多工业应用中的实用解决方案。
最近，硬件算法联合设计框架被广泛研究，以将具有高存储和计算复杂性的深度卷积部署到具有极具有限的存储器（SRAM）和存储器（闪存）的微控制器单元（MCU）。通过搜索最优架构并根据整体网络拓扑安排内存缓冲器，成功地实现了具有大量参数和高FLOPs的深度卷积神经网络部署在MCU，与在高性能计算设备上进行推理的精度相当。视觉Transformer在各种任务中都超过了卷积神经网络，这可以进一步推动资源受限设备中智能视觉感知的极限。
本文中设计了一个硬件算法协同优化框架，以在内存预算极其有限的MCU上部署视觉Transformer。本文共同设计Transformer架构在内存资源约束内实现了最高任务性能。本文共同设计了在内存资源约束瞎实现最高任务性能的Transformer架构，并构建了一个具有高效内存调度的Transformer推理算子库。具体地，利用one-shot网络架构搜索发现最优架构，其中通过考虑低秩分解率和令牌数量来扩大搜索空间，这些低秩分解率和令牌数量显著影响前向传递过程中的峰值内存。为了在大搜索空间中有效地获得有效地SuperNet，本文在由给定低秩分解率和令牌数组成的搜索空间中学习SuperNet，并通过预测任务性能与搜索空间之间相关性来烟花SuperNet优化的搜索空间。为了在内存预算内构建运算符库，本文用整数权重近似替代全精度运算符，如softmax和GeLU。此外，本文将具有大卷积核的补丁嵌入层分解为具有小感受野的多个补丁步骤，并在推理过程中覆盖矩阵乘法中输入令牌特征，以进一步减少内存。
本文方法
有内存约束的网络架构搜索
传统的架构搜索算法通常考虑计算复杂度（FLOPs）与存储成本（参数数量），而峰值内存由于其极其有限的SRAM而掩盖了ViT在微控制器上部署。本文扩展了one-shot NAS以搜索最优架构，其中首先随机学习一些包含所有拓扑选择的超网，并在给定资源预算情况下使用进化算法寻找最优超网。图1展示了具有内存约束的网络架构搜索的框架。为了充分利用ViT NAS的成果，本文在AutoFormer基础上构建了网络的搜索空间，该空间搜索嵌入维度、QKV维度、MLP率、头数量和训练超网的深度。超网的搜索空间是由低秩分解率和可以影响存储器占用令牌数定义的。在推理过程中，令牌数量影响中间特征大小并且可以通过改变补丁嵌入层中感受野来修改。与卷积神经网络的乘累加（MAC）操作不同，大尺寸的矩阵乘法会造成沉重的内存负担。为了解决这个问题，本文将MLPMLP层的低秩分解推广为以下形式：
M
L
P
(
x
)
=
U
r
V
r
T
x
+
b
MLP(x)=U_{r}V_{r}^{T}x+b
MLP(x)=UrVrTx+b
U
r
∈
R
d
×
r
,
V
r
∈
R
d
×
r
U_{r}\in \mathbb{R}^{d\times r},V_{r}\in \mathbb{R^{d\times r}}
Ur∈Rd×r,Vr∈Rd×r是两个估计MLP层原始权重的低秩矩阵。由于低秩矩阵的维度会影响内存占用和任务性能之间的平衡，因此需要在搜索空间选择中找到最优低秩分解化。
在AutoFormer中，将令牌数和低秩分解率直接集成到原始搜索空间中，由于不同低秩分解率的最优架构之间存在较大差异，将令牌数和地址分解率直接集成到原始搜索空间会导致超网的收敛困难。因此首先优化出令牌数和低秩分解率组成的超网的搜索空间，然后再更新后的超网搜索空间中进行架构搜索。重复执行上述两个步骤，直到收敛或达到最大搜索成本。由于给定搜索空间中单个超网的极高训练成本，枚举所有搜索空间对于架构搜索是不可行的。受到Searching the search space of vision transformer启发，由于任务与连续选择之间存在很强的相关性，本文通过考虑性能、内存占用和搜索空间因素（包括令牌数和地址分解率）之间的关系来演化搜索空间。用于评估不同超网搜索空间的最优性的得分函数
S
S
S可以定义为：
S
(
A
,
R
)
=
A
+
η
R
S(A,R)=A+\eta R
S(A,R)=A+ηR
A
A
A是超网在搜索空间的精度，
η
\eta
η是超参数，
R
R
R是内存占用在SRAM限制范围范围内的子网络比率。同时，由于子空间中有大量候选者，枚举超网中所有子网络以获得平均内存是令人望而却步的。使用线性拟合函数
S
=
w
T
x
+
S
0
S=w^{T}x+S_{0}
S=wTx+S0来估计得分和搜索空间的因子之间的趋势。其中
x
=
[
x
1
,
x
2
]
x=[x_{1},x_{2}]
x=[x1,x2]由低秩分解率和令牌数组成。搜索空间将使用以下规则进行更新：
x
i
(
t
+
1
)
=
x
i
(
t
)
−
[
w
i
h
i
]
⋅
Δ
i
x_{i}^{(t+1)}=x_{i}^{(t)}-[\frac{w_{i}}{h_{i}}]\cdot \Delta_{i}
xi(t+1)=xi(t)−[hiwi]⋅Δi
在本文方法中，优化步骤
Δ
i
\Delta_{i}
Δi被设置为相邻因子值的区间。更新阈值
h
i
h_{i}
hi是优化超参数。
在进化的搜索空间中训练超网直到收敛需要不可接受的优化成本，并且拟合了任务性能与训练过程中搜索空间因素的连续选择之间的相关性。对于每个进化的搜索空间，如果不存在预训练的模型，从恢复点或从头训练超网t周期。为了公平地发现任务表现和连续选择之间的趋势，本文将相关权重矩阵
w
w
w与相同总训练周期的数据点进行拟合。由于线性拟合函数仅使用于给定搜索空间的小局部区域，因此使用仅考虑搜索金卡因子
x
x
x的前k最近邻分段线性拟合函数。趋势拟合的数据点采样策略可以表示为：
Φ
(
x
)
=
{
z
∣
z
∈
N
k
(
x
)
,
E
(
z
)
=
E
(
x
)
}
\Phi(x)=\{z|z\in \mathcal{N}_{k}(x),E(z)=E(x)\}
Φ(x)={z∣z∈Nk(x),E(z)=E(x)}
具有k最近邻的
x
x
x领域表示为
N
k
(
x
)
\mathcal{N}_{k}(x)
Nk(x)，并且超网在搜索空间
x
x
x的总训练周期表示为
E
(
x
)
E(x)
E(x)。当给定搜索空间
x
x
x领域为空时，通过考虑总训练周期为
E
(
x
)
−
k
E(x)-k
E(x)−k的前k最近领域来构建领域。本文在这些搜索空间中训练超网t周期以进行精度更新，从而公平地估计给定搜索空间
x
x
x的相关矩阵
w
w
w。
在最优搜索空间中选择的超网用于后续训练，直到完全收敛，并用于ViT的网络架构搜索。在这个过程中，应用进化算法目的是遵守内存约束的同时最大限度提高精度。文章来源：https://www.toymoban.com/news/detail-817226.html
ViT的操作库构建
在微控制器上部署ViT进行推理需要操作库将模型转换为可执行函数。现有的深度学习算子库与CNN相当，但无法在ViT中实现例如GeLU激活函数和层归一化的操作。同时库的推理效率不能充分利用微控制器的内存预算，在较大ViT下实现更高的任务性能。尽管Mcunet: Tiny deep learning on iot devices中高效内存调度技术适用于所有网络架构，但本文提出了以下内存利用率增强技术，以实现在微控制器上部署ViT。
补丁嵌入分解
图像通过ViT的补丁嵌入层进行补丁化进行推理。通过卷积运算将输入样本转换为补丁，其中核大小核步长都等于补丁大小。由于补丁嵌入层中核大小远小于常见的卷积神经网络，因此补丁嵌入层的大部分内存成本来自滤波器权重，而不是输入核输出激活。为了有效解决这一挑战，本文将具有大感受野的卷积转换为具有小感受野的多次卷积运算来分解补丁嵌入层。为了添加多个卷积运算的结果，本文利用一个小的内存缓冲区记录累积的中间值，有效地减少了图像拼接过程中的峰值内存使用量。图2a展示了原始补丁其纳入操作与本文的分解操作之间的比较。由于将补丁嵌入层分解具有小感受野的过量卷积算子会明显增加由于多次前向传递而导致的推理延迟，因此将分解后的补丁嵌入层感受野选择为4x4，以实现内存占用核推理延迟之间令人满意的折衷。
操作融合
量化网络参数核用整数算法代替MAC运算已被证明有利于减少卷积神经网络的内存，而不会导致性能下降。然而，由于包括GeLU激活和层归一化在内的特定算子，具有量化算子的传统算子库对ViT推理是不可行的。由于GeLU需要高斯分布的累计分布函数（CDF），本文用乘法量化和Softmax函数近似激活函数，这两种函数都可以用int8算法实现。类似地，层规范化中平方根运算符不受运算符库中int8支持。使用定点迭代方法构造代理方程来计算平方根运算符的输出。考虑到层归一化算子的推理延迟，只对代理方程迭代3次，并在推理延迟和预测精度之间取得平衡。
令牌重写
ViT推理中的矩阵乘法导致高内存消耗，尽在令牌内施加元素交互，并且令牌之间的交互，并且令牌之间的交互在自注意力计算中启用，其内存成本远低于SRAM限制。因此，可以在全连接层的推理过程中覆盖中间特征，因为在获得相应令牌的输出特征后，将不再使用每个令牌的输入激活。图2b说明了现有推理算子库中矩阵乘法内存调度和本文令牌重写技术，其中本文方法将每个令牌的矩阵乘法输出写入与输入激活相同的内存空间中。由于峰值内存使用发生在矩阵乘法中，根据以下标准，为基于网络架构的ViT推理分配最大内存大小M：
M
=
sup
l
h
f
l
w
f
l
+
max
(
h
i
l
w
i
l
,
h
o
l
w
o
l
)
M=\sup_{l} h_{f}^{l}w_{f}^{l}+\max(h_{i}^{l}w_{i}^{l},h_{o}^{l}w_{o}^{l})
M=lsuphflwfl+max(hilwil,holwol)
h
f
l
,
h
i
l
,
h
o
l
h_{f}^{l},h_{i}^{l},h_{o}^{l}
hfl,hil,hol表示第l全连接层的权重、输入特征和输出特征的高度。
w
f
l
,
w
i
l
,
w
o
l
w_{f}^{l},w_{i}^{l},w_{o}^{l}
wfl,wil,wol表示上述向量的宽度。由于输入激活消耗的内存可以被输出特征覆盖，我们只在输入和输出张量之间分配与较大的内存缓冲区大小相同的内存缓冲。
文章来源地址https://www.toymoban.com/news/detail-817226.html
到了这里，关于MCUFormer: Deploying Vision Transformers on Microcontrollers with Limited Memory的文章就介绍完了。如果您还想了解更多内容，请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章，希望大家以后多多支持TOY模板网！
