用 XNNPACK 加速量化推理

发布人:软件工程师 Marat Dukhan 和 Frank Barchard

量化是在 CPU 上最流行的加速神经网络推理方法之一。去年,TensorFlow Lite 通过 XNNPACK 后端 提高了浮点模型的性能。如今,我们将 XNNPACK 后端扩展至量化模型。各个计算机视觉模型就表现而言,与默认的 TensorFlow Lite 量化内核相比,改扩展在 ARM64 手机上的速度平均提升 30%,在 x86-64 笔记本电脑及桌面设备系统上提升 5 倍,在 WebAssembly SIMD 的浏览器内推理上提升 20 倍。

XNNPACK 中的量化推理针对 TensorFlow 模型优化工具包 所用的 对称量化架构 进行了优化。XNNPACK 既支持 传统的按张量量化架构,也支持较新且准确率更高的 带有权重渠道量化和激活按张量量化的架构。此外,XNNPACK 还支持非对称量化架构,但效率有所降低。

性能提升

我们在一些边缘设备和神经网络架构针对 XNNPACK 加速的量化推理进行了评估。我们以下文介绍的四个公开模型和两个内部量化模型为基准,它涵盖了常见的计算机视觉任务:

1. EfficientNet-Lite0 图像分类 [下载]

2. EfficientDet-Lite0 对象检测 [下载]

3. DeepLab v3 分割,使用 MobileNet v2 特征提取器 [下载]

4. CartoonGAN 图像风格转化 [下载]

5. Face Mesh 特征点的量化版本

6. 视频分割 的量化版本

使用 XNNPACK 在 Android/ARM64 手机上对量化的计算机视觉模型进行单线程推理时的速度提升

在六个 Android ARM64 移动设备上,XNNPACK 与默认的 TensorFlow Lite 量化内核相比,平均提升 30%。

使用 XNNPACK 在 x86-64 笔记本电脑和桌面设备系统上对量化的计算机视觉模型进行单线程推理时的速度提升

XNNPACK 在搭载 x86 处理器的笔记本电脑和桌面设备系统上有了更大幅度的提升。在我们基准测试中的 5 个 x86 处理器上,XNNPACK 将推理速度平均提升 5 倍。值得注意的是,不支持 AVX 指令集的低端和老式处理器,通过将量化推理转换为 XNNPACK,将速度提升 20 倍以上,而 TensorFlow Lite 以前的推理后端只针对 AVX、AVX2 和 AVX512 指令集优化了,而 XNNPACK 则为所有 x86-64 处理器提供了优化实现。

通过 V8 运行时,使用 XNNPACK 在手机、笔记本电脑和桌面设备上对量化的计算机视觉模型进行单线程 WebAssembly SIMD 推理的速度提升

除了传统的移动设备和笔记本电脑/桌面设备平台外,XNNPACK 还通过 TensorFlow Lite Web API 为网络平台进行量化推理加速。上图展示了在 3 个 x86-64 和 2 个 ARM64 系统上通过 V8 JavaScript 引擎运行 WebAssembly SIMD 基准时,与默认的 TensorFlow Lite 实现相比,几何速度平均提升 20 倍。

两年的优化

XNNPACK 起初是 QNNPACK 库的一个分支,但是由于 XNNPACK 的第一个版本专注于浮点推理,而 QNNPACK 专注于量化推理,因此无法将两者进行比较。如今,随着 XNNPACK 引入对量化推理的支持及近两年的性能优化,我们可以直接评估。

为了比较这两个量化推理后端,我们将随机的 MobileNet v1MobileNet v2 模型从 XNNPACK API 移植到 QNNPACK API,并在两个 ARM64 Android 手机和两个 x86-64 系统上对其单线程性能进行了基准测试。上图显示了结果,XNNPACK 在两年内取得了惊人的进步。XNNPACK 在旧的 Pixel 3a 手机上的速度提升 50%,在新的 Pixel 4a 手机上速度提升 4 到 5 倍,在 x86-64 笔记本电脑上提升 2.5X,在 x86-64 工作站上提升 3 倍多。XNNPACK 从 QNNPACK 派生出来后的两年中进行了多种优化,这使其性能提升:

  • XNNPACK 保留了 QNNPACK 中的优化,如 间接卷积算法 和针对微架构的微内核选择,并通过 间接反卷积算法 将其进一步增强,同时具备更灵活的能力,如量化加法运算符和量化乘法运算符中内置的类 numpy 广播。

  • 卷积、反卷积和全连接运算符将 8 位激活和权重的乘积累积为 32 位数字,最终需要将该数字转换回来,或者重新量化为 8 位数字。有多种方法可以实现重新量化,但 QNNPACK 采用了来自 GEMMLOWP库 的架构,该库开创了神经网络推理的量化计算。然而,人们后来发现 GEMMLOWP 的重新量化架构在准确率和性能方面并不是最优的,XNNPACK 凭借更高的性能和准确率取代了它。

  • QNNPACK 以非对称量化架构为目标,其中激活和权重都 表示为无符号整数,并带有零点和比例量化参数,而 XNNPACK 的优化专注于对称量化,其中激活和权重都是有符号整数,且权重有额外的限制:权重的零点始终为零,量化的权重被限制在 [-127, 127] 范围间(虽然 -128 可以表示为有符号 8 位整数,但是仍被排除在外)。在 XNNPACK 中利用对称量化时具有两个计算优势。首先,过滤器权重是静态时,输入零点与过滤器权重的乘积累积结果可以完全融合到卷积、反卷积和全连接运算符中的偏移项中。因此,推理计算中完全没有零点参数。其次,有符号的 8 位输入元素与限制在 [-127, 127] 的权重元素的乘积可以填入 15 位。如此一来,卷积、反卷积和全连接运算符的微内核能够在 16 位变量上进行一半的累积,而不是始终将乘积扩展到 32 位。

  • QNNPACK 微内核针对 ARM 上的 NEON SIMD 指令集和 x86 上的 SSE2 SIMD 指令集进行了优化,但 XNNPACK 支持更多的特定指令集优化。XNNPACK 中的大多数量化微内核都针对 x86/x86-64 上的 SSE2、SSE4.1、AVX、XOP、AVX2 和 AVX512 指令集,ARM/ARM64 上的 NEON、NEON V8 和 NEON 点积指令集,以及 WebAssembly SIMD 指令集进行了优化。此外,XNNPACK 为 WebAssembly 1.0 和 pre-NEON ARM 处理器提供标量支持。

  • QNNPACK 为高端 ARM 和低端 ARM 内核引入了提供专门汇编微内核的想法,而 XNNPACK 将这一想法进一步推进。XNNPACK 不仅针对 Cortex-A53、Cortex-A55 以及具有或不具有 NEON 点积指令集的高端内核的专门专家调整软件流水线汇编微内核,甚至在它们之间还支持实时切换。进行推理的线程从大核迁移到小核时,XNNPACK 会自动适应,从使用针对大核优化的微内核转换到针对小核优化的微内核。

  • QNNPACK 主要专注于多线程推理,并将计算分割成大量的小任务,每个任务计算输出张量的一个小切片。XNNPACK 重新设计了并行化,让任务变得灵活:可以分割为细粒度或粗粒度,具体取决于参与并行化的线程数量。通过动态调整任务粒度,XNNPACK 在单线程执行中实现了低开销,在多线程推理中实现了高并行化效率。

总之,这些优化让 XNNPACK 拥有最先进的量化推理能力,且让 TensorFlow Lite 成为最通用的量化推理解决方案,涵盖包括 Raspberry Pi Zero、Chromebook,以及具有服务器类处理器的工作站在内的众多系统。

如何使用?

量化的 XNNPACK 推理在所有平台 TensorFlow Lite 的 CMake 构建,以及网络平台 TensorFlow Lite 的 Bazel 构建中默认启用,并将在 2.7 版本的 TensorFlow Lite Web API 中可用。在其他平台的 Bazel 构建中,量化的 XNNPACK 推理通过构建时的选择机制启用。使用 Bazel 构建 TensorFlow Lite 时,如果添加 --define tflite_with_xnnpack=true --define xnn_enable_qs8=true,TensorFlow Lite 解释器将默认使用 XNNPACK 后端,使用具有对称量化的受支持运算符。可以通过添加 --define xnn_enable_qu8=true Bazel 选项,对非对称量化的运算符进行有限支持。

可加速哪些运算?

XNNPACK 后端目前支持量化的 TensorFlow Lite 运算符的子集(请参阅 文档,了解详细信息和限制)。XNNPACK 支持由模型优化工具包通过训练后的整数量化和量化感知训练产生的模型,但不支持训练后的动态范围量化。

未来展望

这是 XNNPACK 集成到 TensorFlow Lite 后的第三个版本,之前是 浮点实现的初始版本 以及 引入稀疏推理支持的后续版本。在未来版本中,我们将添加以下改进:

  • 最新 ARM 处理器上的半精度推理
  • 稀疏量化推理。
  • 更加快速的密集推理。

希望您能在 GitHubStackOverflow 页面上积极发表您的想法和评论,您也可以在 discuss.tensorflow.org 上提问。

原文:Faster Quantized Inference with XNNPACK
中文:TensorFlow 公众号