使用 XNNPACK 集成加速 TensorFlow Lite
2020 年 7 月 24 日
作者:Marat Dukhan,Google Research

利用 CPU 进行 ML 推理可以实现最广泛的边缘设备覆盖范围。因此,提高 CPU 上的神经网络推理性能一直是 TensorFlow Lite 团队的首要要求之一。我们倾听了您的意见,并很高兴地向您提供平均 2.3 倍更快的浮点推理速度,这是通过将 XNNPACK 库集成到 TensorFlow Lite 中实现的。

为了实现这种加速,XNNPACK 库 提供了高度优化的浮点神经网络运算符实现。它在今年早些时候推出了 TensorFlow.js 的 WebAssembly 后端,并且在这个版本中,我们引入了针对 TensorFlow Lite 使用案例的额外优化。
  • 为了为移动设备上的 TensorFlow Lite 用户提供最佳性能,所有运算符都针对 ARM NEON 进行了优化。最关键的运算符(卷积、深度可分离卷积、转置卷积、全连接)针对移动电话中常用的 ARM 内核(例如 Pixel 2 中的 Cortex-A53/A73 和 Pixel 3 中的 Cortex-A55/A75)进行了汇编优化。
  • 对于 x86-64 设备上的 TensorFlow Lite 用户,XNNPACK 添加了对 SSE2、SSE4、AVX、AVX2 和 AVX512 指令集的优化。
  • XNNPACK 不是逐个执行 TensorFlow Lite 运算符,而是查看整个计算图并通过运算符融合对其进行优化。例如,在 TensorFlow Lite 中,带有显式填充的卷积由 PAD 运算符和具有 VALID 填充模式的 CONV_2D 运算符的组合表示。XNNPACK 检测到此运算符组合并将这两个运算符融合成一个带有显式指定填充的卷积运算符。
用于 CPU 的 XNNPACK 后端加入了 TensorFlow Lite 加速推理引擎家族,用于 移动 GPUAndroid 的神经网络 APIHexagon DSP边缘 TPUApple 神经引擎。它提供了一个强大的基线,可以在所有移动设备、桌面系统和树莓派板上使用。

在 TensorFlow 2.3 版本中,XNNPACK 后端包含在 Android 和 iOS 的预构建 TensorFlow Lite 二进制文件中,并且可以通过一行代码更改启用。Windows、macOS 和 Linux 构建的 TensorFlow Lite 也支持 XNNPACK 后端,它通过构建时选择加入机制启用。在更广泛的测试和社区反馈之后,我们计划在即将发布的版本中将其默认启用到所有平台。

性能改进

TensorFlow Lite 中的 XNNPACK 加速推理已在 Google 产品中投入生产,我们观察到各种神经网络架构和移动处理器都实现了显著的加速。XNNPACK 后端将 Pixel 3a Playground 中的背景分割速度提高了 5 倍,并将 ARCore 中 增强面部 API 中的神经网络模型速度提高了 2 倍。

我们发现,TensorFlow Lite 从小型神经网络模型和低端移动电话的 XNNPACK 后端中获益最大。下面,我们展示了针对 9 个公共模型的基准测试,这些模型涵盖了常见的计算机视觉任务。

  1. MobileNet v2 图像分类 [下载]
  2. MobileNet v3-Small 图像分类 [下载]
  3. DeepLab v3 分割 [下载]
  4. BlazeFace 人脸检测 [下载]
  5. SSDLite 2D 对象检测 [下载]
  6. Objectron 3D 对象检测 [下载]
  7. Face Mesh 关键点 [下载]
  8. MediaPipe Hands 关键点 [下载]
  9. KNIFT 局部特征描述符 [下载]
使用 XNNPACK 后端的 TensorFlow Lite 与默认后端相比的单线程推理加速,跨越 5 款移动电话。数字越大越好。
使用 XNNPACK 后端的 TensorFlow Lite 与默认后端相比的单线程推理加速,跨越 5 台台式机、笔记本电脑和嵌入式设备。数字越大越好。

如何使用它?

XNNPACK 后端已包含在预构建的 TensorFlow Lite 2.3 二进制文件中,但需要显式运行时选择加入才能启用它。我们正在努力在将来的版本中默认启用它。

在 Android/Java 上选择加入 XNNPACK 后端

预构建的 TensorFlow Lite 2.3 Android 存档 (AAR) 已包含 XNNPACK,只需一行代码即可在 Interpreter.Options 对象中启用它。
Interpreter.Options interpreterOptions = new Interpreter.Options();
interpreterOptions.setUseXNNPACK(true);
Interpreter interpreter = new Interpreter(model, interpreterOptions);

在 iOS/Swift 上选择加入 XNNPACK 后端

预构建的 TensorFlow Lite 2.3 CocoaPods for iOS 同样包含 XNNPACK,以及在 InterpreterOptions 类中启用它的机制。
var options = InterpreterOptions()
options.isXNNPackEnabled = true
var interpreter = try Interpreter(modelPath: "model/path", options: options)

在 iOS/Objective-C 上选择加入 XNNPACK 后端

在 iOS 上,可以通过 TFLInterpreterOptions 类中的新属性从 Objective-C 启用 XNNPACK 推理。
TFLInterpreterOptions *options = [[TFLInterpreterOptions alloc] init];
options.useXNNPACK = YES;
NSError *error;
TFLInterpreter *interpreter =
    [[TFLInterpreter alloc] initWithModelPath:@"model/path"
                                      options:options
                                        error:&error];

在 Windows、Linux 和 Mac 上选择加入 XNNPACK 后端

在 Windows、Linux 和 Mac 上,XNNPACK 后端通过构建时选择加入机制启用。使用 Bazel 构建 TensorFlow Lite 时,只需添加 --define tflite_with_xnnpack=true,TensorFlow Lite 解释器将默认使用 XNNPACK 后端。

使用您的 TensorFlow Lite 模型试用 XNNPACK

您可以使用 TensorFlow Lite 的 基准测试工具 并在 XNNPACK 中测量 TensorFlow Lite 模型的性能。您只需使用 --use_xnnpack=true 标志启用 XNNPACK,即使基准测试工具是在没有 --define tflite_with_xnnpack=true Bazel 选项的情况下构建的。
adb shell /data/local/tmp/benchmark_model \
  --graph=/data/local/tmp/mobilenet_quant_v1_224.tflite \
  --use_xnnpack=true \
  --num_threads=4

哪些操作已加速?

XNNPACK 后端目前支持一部分浮点 TensorFlow Lite 运算符(有关详细信息和限制,请参阅 文档)。XNNPACK 支持 32 位浮点模型和使用 16 位浮点量化权重的模型,但不支持权重或激活中使用定点量化的模型。但是,您不必将模型限制为 XNNPACK 支持的运算符:任何不受支持的运算符将透明地回退到 TensorFlow Lite 中的默认实现。

未来工作

这只是 XNNPACK 后端的第一个版本。除了社区反馈之外,我们还打算添加以下改进。
  • 集成 快速稀疏卷积网络 算法
  • 在最新的 ARM 处理器上进行半精度推理
  • 在定点表示中进行量化推理

我们鼓励您在我们的 GitHubStackOverflow 页面上留下您的想法和评论。

致谢

我们要感谢 Frank Barchard、Chao Mei、Erich Elsen、Yunlu Li、Jared Duke、Artsiom Ablavatski、Juhyun Lee、Andrei Kulik、Matthias Grundmann、Sameer Agarwal、Ming Guang Yong、Lawrence Chan、Sarah Sirajuddin 的贡献。
下一篇文章
Accelerating TensorFlow Lite with XNNPACK Integration

作者:Marat Dukhan,Google Research

利用 CPU 进行 ML 推理可以实现最广泛的边缘设备覆盖范围。因此,提高 CPU 上的神经网络推理性能一直是 TensorFlow Lite 团队的首要要求之一。我们倾听了您的意见,并很高兴地向您提供平均 2.3 倍更快的浮点推理速度,这是通过将 XNNPACK 库集成到 TensorFlow Lite 中实现的。…