ITPub博客

首页 > Linux操作系统 > Linux操作系统 > TensorFlow Lite CPU 推理提速四倍:XNNPack 动态范围量化技术解析

TensorFlow Lite CPU 推理提速四倍:XNNPack 动态范围量化技术解析

Linux操作系统 作者:lumin1020 时间:2026-09-15 12:12:45 0 删除 编辑

导语

2024 年 4 月,TensorFlow 官方博客发布了一项重要更新:XNNPack 的全连接和二维卷积算子现已支持动态范围量化。这一改进让 TensorFlow Lite 的 CPU 推理性能相比单精度基线提升了四倍,意味着更多 AI 功能可以部署到更旧、更低端的设备上。对于长期受限于端侧算力与模型精度平衡难题的开发者而言,这无疑是一个值得关注的进展。

背景:CPU 仍是端侧推理的主战场

XNNPack 是 TensorFlow Lite 的 CPU 后端。尽管 GPU、NPU 等专用加速器在移动端日益普及,CPU 依然为机器学习推理提供了最广泛的覆盖范围,并且仍然是 TensorFlow Lite 的默认目标。因此,提升 CPU 推理性能始终是 TensorFlow Lite 团队的首要任务。

在此之前,XNNPack 为用户提供的选择相对有限:要么使用完全整数量化(权重和激活值均存储为有符号 8 位整数),要么使用半精度(fp16)或单精度(fp32)浮点推理。完全整数量化虽然性能优异,但转换过程复杂、容易出错,且精度难以保证;浮点推理虽然精度高,但在低端设备上性能捉襟见肘。动态范围量化的引入,正是为了在这两者之间找到平衡点。

动态范围量化是什么?

动态量化模型与完全量化模型有相似之处:全连接和卷积算子的权重在模型转换期间被量化为 8 位整数。但关键区别在于,所有其他张量不被量化,保持为 float32 张量。

在模型推理期间,浮点层激活值在传递给全连接和卷积算子之前被转换为 8 位整数。激活张量每一行的量化参数(零点 zero point 和缩放因子 scale)根据观测到的激活值范围动态计算。这最大化了量化过程的精度,因为激活值充分利用了 8 个量化位。

相比之下,在完全量化模型中,这些参数在模型转换期间基于使用代表性数据集观测到的激活值范围固定下来。此外,完全量化算子的输出为 8 位整数,而动态范围量化算子的输出为 32 位浮点格式。

通过动态范围量化,开发者获得了完全量化的大部分性能收益,同时保持了更高的整体精度。

性能表现:数据说话

TensorFlow 团队在四个公共模型上进行了基准测试,涵盖常见的计算机视觉任务:

  1. EfficientNetV2 — 图像分类和特征提取
  2. Inception-v3 — 图像分类
  3. Deeplab-v3 — 语义分割
  4. Stable Diffusion — 图像生成(扩散模型)

每个模型在可能的情况下进行了三种转换:全浮点、全 8 位有符号整数量化和动态范围量化。Stable Diffusion 的扩散模型由于不支持的操作符,无法使用全整数量化进行转换。

测试结果令人瞩目:

  • 动态范围量化与全整数量化相比具有竞争力,在某些情况下甚至可以超过其性能。
  • Stable Diffusion 的扩散模型运行速度比原始 float32 模型快高达 6.2 倍。
  • 在测试的三个模型中,有两个模型的全整数量化并未如预期般快于动态范围量化。

对于这一反直觉的结果,团队使用 TFLite 的分析器进行模型分析后解开了谜团:这些模型较慢是由于量化伪影的组合——当输入和输出尺度的比率落在某个范围内时,量化算术更高效,以及 XNNPack 中缺少算子支持。由于尺度的比率落在最优范围之外,必须采取次优路径,性能因此受到影响。

混合精度推理:动态量化 + fp16

半精度和动态范围量化现在可以在 XNNPack 中结合使用,以在具有硬件 fp16 支持的设备上获得最佳的设备端 CPU 推理性能。全连接和二维卷积算子可以输出 fp16 数据而非 fp32。

与 fp32 相比,fp16 使用一半的位数来存储浮点值,这意味着由于尾数显著缩短(10 位 vs 23 位),每个值的相对精度会降低。但并非所有模型都支持 fp16 推理,如果模型支持,向量化浮点算子的计算成本可以减半,因为 CPU 每条指令可以处理两倍的数据。具有计算密集型浮点算子(如批量矩阵乘法和 Softmax)的动态量化模型也可以受益于 fp16 推理。

团队通过对比实验验证了混合精度推理的可行性:使用 fp32 激活的动态量化 Stable Diffusion 模型与使用 fp16 激活生成的图像,在相同随机数种子下无法区分。这有力地表明扩散模型适合 fp16 推理。当然,与所有神经网络一样,任何量化策略都应使用大型验证数据集进行验证,而不仅仅是单次测试。

如何使用动态范围量化?

使用动态范围量化需要两个步骤:

第一步:模型转换。 从 TensorFlow 转换模型并支持动态范围量化。已经使用动态范围量化转换过的现有模型无需重新转换。在模型转换期间,可以通过启用 converter.optimizations = [tf.lite.Optimize.DEFAULT] 转换器标志来启用动态范围量化。

与完全整数量化不同,动态范围量化不需要代表性数据集,且不支持的算子不会阻止转换成功。因此,动态范围量化比完全整数量化对非专业用户来说更加易于使用。

第二步:推理部署。 从 TensorFlow 2.17 开始,动态量化的 XNNPack 推理将在预构建的二进制文件中默认启用。如果想更早使用,可以使用每夜构建的 TensorFlow 版本。

架构支持与落地情况

这些算子针对 XNNPack 支持的所有架构进行了优化,包括 ARM、ARM64、x86 SSE/AVX/AVX512 和 WebAssembly,涵盖最新的 ArmV9 处理器,如 Pixel 8 的 Tensor G3 CPU 或 One Plus 11 的 SnapDragon 8 Gen 2 CPU。

在落地方面,XNNPack 的动态范围量化现已应用于 Gemini、Google Meet 和 Chrome OS 音频降噪,并将在今年登陆更多产品。这项技术现已向开源用户开放。

趋势判断:端侧推理的务实路线

从这次更新中,可以读出几个值得关注的趋势:

其一,动态范围量化正在成为端侧推理的“甜点”方案。 完全整数量化虽然理论性能最优,但转换门槛高、精度风险大;纯浮点推理则难以满足低端设备的性能要求。动态范围量化在两者之间提供了一种折中:模型大小与全量化模型相近,性能提升通常相似,有时甚至超过全量化模型,且转换过程大大简化。

其二,混合精度策略正在走向成熟。 动态量化与 fp16 的结合,让开发者可以根据设备能力灵活选择激活精度。在支持 fp16 的设备上获得最佳性能,在不支持的设备上回退到 fp32,这种分层策略有望成为端侧推理的标配。

其三,CPU 推理的优化远未触及天花板。 通过算子层面的持续优化(如 XNNPack 对全连接和卷积算子的逐架构优化),CPU 在端侧推理中的角色依然不可替代。对于需要广泛设备覆盖的应用场景,CPU 后端仍是最稳妥的选择。

结论

XNNPack 对动态范围量化的支持,为 TensorFlow Lite 开发者提供了一个兼顾性能与易用性的新选项。四倍的推理性能提升、无需代表性数据集的转换流程、以及与 fp16 混合精度的兼容性,都让这一技术具备了大规模落地的条件。随着 TensorFlow 2.17 的发布,动态量化 XNNPack 推理将默认启用,端侧 AI 应用的性能门槛有望进一步降低。

对于正在探索端侧模型部署的团队而言,依据现有素材判断,动态范围量化值得纳入技术选型的评估范围。建议结合自身模型结构和目标设备架构,使用大型验证数据集进行充分的精度与性能验证。


生成标记:本文由 ITPUB 文章生成平台基于已选素材整理生成。

来自 “ ITPUB博客 ” ,链接:https://blog.itpub.net/69965512/viewspace-3182408/ ,如需转载,请注明出处,否则将追究法律责任。

请登录后发表评论 登录
全部评论

注册时间:2020-03-06

  • 博文量
    299
  • 访问量
    473949