/ ai资讯

利用ExecuTorch和Arm SME2加速端侧机器学习推理

发布时间:2026-03-03 11:46:32

作者:Arm 高级首席工程师 Jason Zhu 等人

交互式图像分割已成为全球主流应用中极具代表性的移动端体验。简单来说,用户只需在图片上轻点一下(或粗略勾画),应用就能立刻生成像素蒙版,把目标对象“抠”出来。这项技术支撑了许多常见功能,比如制作个性化贴纸、分离主体以替换背景,或是对图像局部进行选择性增强。这些效果背后,是轻量级分割模型在运行,这些模型通过 ExecuTorch(PyTorch 的开源端侧推理运行时)以及第二代 Arm 可伸缩矩阵扩展技术 (Arm SME2) 运行。

本文将探讨这些软硬件技术升级如何让抠图功能背后的端侧交互式分割模型 SqueezeSAM 在图像分割任务中实现最高可达 3.9 倍的加速,并阐述这一突破对移动端应用开发者的广泛影响。SqueezeSAM 已部署在 Meta 旗下应用中。

移动设备上端侧 AI 的兴起

随着端侧人工智能 (AI) 不断发展,一个核心问题摆在眼前:当更强大的模型在严格的移动端功耗与时延限制下能够运行得更快时,会出现哪些新的可能性?实际上,许多交互式移动端 AI 功能和工作负载已在 CPU 上运行,因为 CPU 始终可用、与应用无缝集成,且在各类场景中具备高灵活性、低时延与出色性能。对于这类部署方案,性能优劣往往取决于 CPU 执行矩阵密集型内核的效率,以及当算力不再是瓶颈后还存在哪些限制因素。

SME2 是 Armv9 架构中的一组高级 CPU 指令,专为在端侧直接加速面向矩阵的计算工作负载而设计。我们量化了在 ExecuTorch 与 XNNPACK 部署方案中,SME2 对端到端推理的加速效果,并通过算子级性能分析展示具体哪些方面得到了改进。启用 SME2 的全新 Arm CPU已集成在Arm Lumex 计算子系统 (CSS)中,用于旗舰智能手机与下一代 PC 设备。

案例研究:利用 SME2 加速交互式图像分割

我们评测了在以 ExecuTorch 和 XNNPACK 为后端运行时,SME2 对 SqueezeSAM 推理时延的影响。该方案利用 Arm KleidiAI 优化的内核,以发挥 SME2 的加速能力。

启用 SME2 后,8 位整型 (INT8) 和 16 位浮点型 (FP16) 推理均获得显著性能提升(图 1)。在采用默认功耗配置的单个 CPU 核心上,INT8 时延优化 1.83 倍(从 556 毫秒降至 304 毫秒),FP16 时延优化 3.9 倍(从 1163 毫秒降至 298 毫秒)。若无 SME2,时延会过高,无法满足交互式场景的流畅使用需求;启用 SME2 后,单核端到端推理时延可达到 300 毫秒左右,使端侧部署切实可行,同时也为应用的其他部分留出了性能余量。

上述结果表明,SME2 可在 CPU 上显著加速量化后的 INT8 模型。同时,在本案例中,SME2 让 FP16 时延接近 INT8 水平,这一成果意义重大,因为它并非替代 INT8,而是扩展了实际可部署的方案范围。这让开发者拥有更高的灵活性,可选择最符合精度与工作流需求的数据格式,尤其适用于对精度敏感的工作负载,如图像超分辨率、图像抠图、暗光去噪与高动态范围 (HDR) 增强。倘若没有如此级别的 FP16 加速,移动端部署通常只能选用 INT8 以满足时延目标,而这意味着需要引入量化工作流并承担精度下降的风险。

除基准测试数据外,这些性能提升可直接转化为可用的 CPU 算力余量。这些余量可用于打造更丰富的体验,例如在保持相机预览与 UI 流畅的同时,并行运行分割与增强任务(如去噪或 HDR 处理);或者把原本只能处理单张图片的抠图,扩展成带跨帧目标跟踪的实时视频抠图,亦可用于降低功耗。

图 1:普通模式下(默认移动端功耗设置),一个 CPU 核心在启用与禁用 SME2 时 SqueezeSAM 的端到端时延。INT8 从 556 毫秒优化至 304 毫秒(提升 1.83 倍)。FP16 从 1163 毫秒优化至 298 毫秒(提升 3.90 倍),在本案例中 FP16 时延已接近 INT8 水平。

本文所有结果均为在搭载启用 SME2 的 Arm CPU 的旗舰安卓智能手机上进行受控测试所得。性能会因模型、硬件及具体设备设置而异。

技术栈:PyTorch、ExecuTorch、XNNPACK、KleidiAI 和 SME2

框架间的连接关系

上图总结了本案例研究中使用的 CPU 执行技术栈。模型在 PyTorch 中定义,由 ExecuTorch 导出并运行,CPU 计算则委派给作为后端的 XNNPACK 执行。XNNPACK 使用 Arm KleidiAI,这是面向 Arm CPU、为加速机器学习工作负载而优化的轻量级 CPU 内核库。这些内核可在受支持的设备上自动利用 SME2 加速,同时也能为不支持 SME2 的系统提供针对其他的 CPU 特性的优化实现。

当 ExecuTorch 在启用 XNNPACK 委派的情况下运行模型时,XNNPACK 会在运行时根据底层硬件能力选择合适的内核实现。在启用 SME2 的设备上,这些运算中的矩阵乘法计算可直接受益于 SME2 加速,无需对模型结构或应用代码进行任何修改。在这类运算得到加速后,推理管线中的其他环节(如数据移动、布局转换、未委派的算子等)往往会成为新的性能瓶颈。这也是算子级性能分析对于理解端到端性能至关重要的原因。

案例研究模型

在本次评估中,我们使用了 SqueezeSAM 模型,该模型采用轻量级、以 conv2d 为主的 UNet 架构,是典型的移动端视觉模型。

模型结构可被映射为两大类工作,这两类工作对端到端推理时间有着显著影响:

计算密集型运算:卷积层(iGEMM,隐式通用矩阵乘法)和注意力/MLP 层(GEMM,通用矩阵乘法)。

数据移动类运算:转置、维度重塑和布局转换。

平台说明:在许多基于 Armv9 架构的设备上,SME2 作为 CPU 核心间的共享执行资源实现,其伸缩特性会随系统级芯片 (SoC) 与 CPU 微架构不同而存在差异。我们在评估中已明确考虑这一点,并在解读单核与多核结果时讨论其产生的影响。

结果:INT8 和 FP16

(1 个 CPU 核心对比 4 个 CPU 核心)

我们在启用与禁用 SME2 的条件下,对同一模型的两种精度(INT8 和 FP16)进行基准测试。我们重点关注单核执行场景(SME2 在此场景下相对收益最大),同时也给出四核结果,以说明当 SME2 作为共享硬件资源时的绝对时延与伸缩表现。所有测试均仅统计模型本身的推理时延。

模型通过 ExecuTorch 在安卓智能手机上运行,在相同软件与系统环境下分别测试启用与禁用 SME2 的情况。除非另有说明,所有结果均为在无温控降频情况下的稳态性能。

所有结果均以“普通模式 | 无限制模式(毫秒)”的形式给出。普通模式对应默认的移动设备电源设置,即系统电源策略启用状态,反映典型用户使用场景。无限制模式对应持续供电、保持唤醒状态的配置,CPU 频率限制有效解除;单核测试中,无限制模式结果固定在最高性能(Ultra/Prime,本例中为 4.2 GHz)CPU 核心上运行。

在两种模式下,SME2 均呈现一致的相对加速趋势,表明尽管绝对时延存在差异,但其加速效果受系统功耗策略影响较小。除非另有明确说明,本文后续均以普通模式结果为主,因其更能反映典型手机使用环境下的用户感知时延。无限制模式结果用于展示性能余量与硬件上限,应视为最佳表现,而非日常用户体验。

表 1:在安卓手机上启用与禁用 SME2 时,SqueezeSAM 的端到端时延结果,分别在一个 CPU 核心与四个 CPU 核心上测试(仅模型时延)。数值以“普通模式 | 无限制模式(毫秒)”的形式给出。

关于四核扩展说明:四核的加速比例较小(例如,普通模式下 INT8 为 1.08 倍,而单核为 1.83 倍),这与 SME2 作为共享资源的特性一致,同时也受内存带宽、缓存行为等其他系统共享因素影响。伸缩特性会因 SoC 与 CPU 实现方式不同而存在差异。在生产部署中,若能满足时延目标,优先使用一到两个核心可获得更好的能效;当需要更低的绝对时延且功耗预算允许时,可使用更多核心。

算子级性能分析的重要性

端到端时延只能告诉我们性能提升了多少,无法说明原因及后续的优化对象。为了理解 SME2 的性能收益来源及下一阶段的性能瓶颈,我们使用算子级性能分析。

我们通过 ExecuTorch 开发工具中的性能分析工具 ETDump 采集每个算子的耗时信息,该工具会记录推理过程中各个算子的执行时间。这使我们能够将端到端加速效果归因到模型的具体部分,如图 2 和表 2 所示。

为了让分析更具实践指导意义,我们将算子归纳为少数几个与常见模型结构精准对应的类别:

卷积:Conv2d 层(通常基于 iGEMM 实现);

GEMM:矩阵乘法和线性层(注意力和 MLP 投影);

逐元素:ReLU、GELU、Add、Mul 及其他逐点运算;

数据移动:转置、拷贝、格式转换、维度重塑和填充;

其他:未委派的算子和框架开销。

通过上述分类拆解,我们可以明确 SME2 在哪些方面作用最为显著,以及在矩阵计算被加速后依然存在的性能瓶颈。

图 2:在安卓智能手机上(一个 Arm CPU 核心,默认移动设备功耗设置),启用与禁用 SME2 时,FP16 与 INT8 的算子类别耗时明细(绝对时间)。SME2 大幅降低卷积与 GEMM 耗时,数据移动在运行时间中的占比显著提升。

表 2:在安卓智能手机上(一个 Arm CPU 核心,默认移动设备功耗设置),禁用与启用 SME2 情况下 INT8 与 FP16 的算子级耗时明细对比。非矩阵乘法算子主要受运行时波动的影响。

从端到端与算子级结果得出的三大洞察

洞察 1:SME2 能够加速矩阵计算,将瓶颈转移至数据移动

SME2 显著降低 INT8 与 FP16 精度下的端到端时延。在单个 Arm CPU 核心上,INT8 性能优化 1.83 倍(从 556 毫秒降至 304 毫秒),FP16 优化 3.90 倍(从 1163 毫秒降至 298 毫秒)。即使在四核场景下,SME2 仍可大幅降低 FP16 时延(从 374 毫秒降至 193 毫秒)。这些优化效果使单核执行时延进入约 300 毫秒区间,在为应用其他部分保留 CPU 余量的同时,让端侧实时交互成为可能。

算子级性能分析表明,SME2 能够大幅加速矩阵密集型算子。禁用 SME2 时,卷积与 GEMM 占据推理的主要耗时,分别占 INT8 运行时间的 55.7%、FP16 的 75.8%。启用 SME2 后,GEMM 算子加速约 3 至 4 倍,卷积/iGEMM 加速约 4 至 9 倍,这是端到端性能提升的主要驱动因素。

矩阵计算加速后,数据移动与框架开销的相对占比上升,后续优化重心也随之转移。

洞察 2:由转置驱动的数据移动约占总运行时的 40%

在 SME2 加速后,数据移动成为主要运行耗时因素之一。在启用 SME2 的 INT8 运行中,数据移动占总运行时的 41.4%(FP16 为 39.9%)。ETDump 追踪结果显示,约 85% 的数据移动时间来自转置算子,仅两类转置节点就占用了该类别超过 80% 的耗时。

这类开销源于模型不同部分与运行时之间的数据布局不匹配,而非计算强度问题。实际场景中,当具有不同布局偏好的算子按序组合时,会触发频繁的 NCHW NHWC 格式转换。在本模型中可以看到:归一化算子作为可移植的 NCHW 算子执行,且无法与相邻卷积融合(例如当非线性激活函数位于 Conv2d 与 BatchNorm 之间时),而 XNNPACK 卷积内核更偏好 NHWC 布局。这会在 UNet 编码器–解码器模块中引发重复的布局转换:

BatchNorm/GroupNorm (NCHW) →

转置 (NCHW→NHWC) → 卷积 (NHWC) →

转置 (NHWC→NCHW) →

BatchNorm/GroupNorm (NCHW)

由于这类开销由模型与运行时的布局选择决定,而非计算强度,因此必须通过性能分析才能将其暴露出来,从而将其转化为可执行的优化目标。

重要的是,这一性能分析洞察已被证实具备实际优化价值。作为初步举措,Meta ExecuTorch 团队在框架中实现了针对性的图级优化,以减少归一化层周围不必要的数据布局转换。在我们的实验中,除 SME2 带来的加速收益外,还可使 INT8 时延额外减少约 70 毫秒 (23%),FP16 时延额外减少约 30 毫秒 (10%)。

由上述结果可以确认,高转置的数据移动是极具价值的优化方向。随着我们持续分析整张计算图的布局行为,仍有进一步优化的潜力。

洞察 3:在本案例研究中,启用 SME2 后 FP16 时延接近 INT8 水平

尽管 INT8 每个张量元素仅占用一半的内存带宽,但这并不直接带来成比例的端到端加速。启用 SME2 后,本案例中 FP16 时延已接近 INT8(单个核心上分别为 298 毫秒与 304 毫秒)。

算子耗时明细揭示了背后原因。FP16 的卷积加速效果尤为显著(加速 9.0 倍,INT8 为 4.4 倍),弥补了 INT8 在内存上的效率优势。同时,INT8 矩阵计算路径会带来额外开销,包括量化、伸缩及更复杂的内核调度逻辑,削弱了 INT8 的有效带宽优势。

最终效果是,SME2 拓宽了可选用的精度范围。INT8 依然是高效方案,而对于不希望承担量化复杂度或精度损耗的精度敏感型工作负载,FP16 也变得更加实用。尽管本案例中 FP16 性能已接近 INT8,但该效果与任务负载强相关,会随算子组合、张量形状与内存压力发生变化。

实操示例:重现工作流

如想自行尝试上述工作流,我们提供了基于开源 SAM 模型的实操教程,内容涵盖模型导出、使用 SME2 执行推理、通过 ETDump 进行算子级性能分析等。完整的设置说明与代码示例可在代码仓库及 Learning Paths 中获取。

代码仓库: https://github.com/ArmDeveloperEcosystem/sme-executorch-profiling

Learning Paths: https://learn.arm.com/learning-paths/cross-platform/sme-executorch-profiling/

你将能学到什么:

如何将分割模型导出至 ExecuTorch,并启用 XNNPACK 委派

如何在已启用 SME2 的安卓、iOS 和 macOS 设备上构建与部署模型

如何运行 ETDump 性能分析,采集各算子的耗时信息

如何在自有模型中识别并量化数据移动及其他非计算类性能瓶颈

结论:SME2 带来的实际改变

在本 SqueezeSAM 案例研究中,SME2 为 INT8 与 FP16 提供了显著的端侧 CPU 加速效果,从本质上提升了交互式移动端工作负载的可行性。

这对开发者和产品团队意味着什么:

端侧机器学习在 CPU 上更具可行性:SME2 可实现最高 3.9 倍的端到端推理加速。在安卓默认功耗设置下,真实交互式移动端模型的单核时延可从 1 秒以上降至约 300 毫秒。对于交互式工作负载,这使得基于 CPU 的端侧机器学习从勉强可用变为稳定实用,同时为应用其他功能保留性能空间。

FP16 在部分场景中成为更可行的部署选择:SME2 大幅加速 FP16 计算,并缩小其与 INT8 之间的时延差距,让开发者能更灵活地选择最符合精度、工作流与时延要求的数值精度,尤其适用于对精度敏感的工作负载。

节省的算力余量可带来更丰富的使用体验:释放的 CPU 预算可用于增强端侧功能,例如在图像分割的同时运行画质增强(如去噪、HDR),或将图像抠图从单张图片扩展至支持跨帧目标跟踪的实时视频。

性能分析给出下一阶段优化目标:当 SME2 加速了矩阵密集型算子(卷积/iGEMM、GEMM)后,性能瓶颈通常会转向数据移动与未委派算子。基于 ETDump 的算子级性能分析可清晰展示这类开销,并提供可落地的优化方向。

根据起点不同,有两点明确的启示:

若你目前尚未部署端侧机器学习,那么基于 SME2 的 CPU 加速可以让移动端 CPU 成为部署这类“重算子”模型的可行起点,而性能分析能够为验证表现和持续迭代提供清晰路径。

若你已部署端侧模型,SME2 可释放算力余量,用于拓展功能、提升用户体验;同时性能分析可指出收益最高的后续优化方向(在 SqueezeSAM 中,由转置驱动的布局转换约占总运行时间的 40%)。

综上,SME2 加速与算子级性能分析相结合,可形成一套实用工作流:既能快速获得立竿见影的性能提升,亦可精准定位端侧 AI 后续的重点优化方向。

  • ARM ARM 关注

    关注

    135

    文章

    9552

    浏览量

    391756

免责声明:本文为转载,非本网原创内容,不代表本网观点。其原创性以及文中陈述文字和内容未经本站证实,对本文以及其中全部或者部分内容、文字的真实性、完整性、及时性本站不作任何保证或承诺,请读者仅作参考,并请自行核实相关内容。

如有疑问请发送邮件至:bangqikeconnect@gmail.com