作者:PyTorch 团队

今天,我们宣布推出四项 PyTorch 原型特性。其中前三项将使移动机器学习开发者能够在使用片上系统 (SOC) 的全部硬件 (HW) 引擎上执行模型。这为开发者提供了优化模型执行的选项,以实现独特的性能、功耗和系统级并发性。

这些特性包括支持在以下设备上硬件引擎上执行:

  • 利用 Android 神经网络 API (NNAPI) 支持 DSP 和 NPU 执行,这是与 Google 合作开发的
  • 通过 Vulkan 在 Android 上执行 GPU 计算
  • 通过 Metal 在 iOS 上执行 GPU 计算

本次发布还通过新增的 Linux ARM64 构建支持带来了开发者效率方面的优势。

下面是各项特性的简要说明以及入门链接。这些特性可通过我们的每夜构建版本获取。欢迎通过PyTorch 论坛与我们联系并提供任何评论或反馈。我们非常期待收到您的反馈,并了解您如何使用这些特性!

支持 Google Android NNAPI

Google Android 和 PyTorch 团队合作,通过 PyTorch Mobile 支持 Android 的神经网络 API (NNAPI)。开发者现在可以在 Android 手机上实现高性能执行,因为他们的机器学习模型将能够访问手机片上系统上的额外硬件模块。NNAPI 允许 Android 应用在为移动手机提供动力的芯片上最强大和最高效的部分运行计算密集型神经网络,包括 DSP(数字信号处理器)和 NPU(专用神经网络处理器)。该 API 在 Android 8 (Oreo) 中引入,并在 Android 10 和 11 中进行了显著扩展,以支持更丰富的 AI 模型集。通过此次集成,开发者现在可以直接从 PyTorch Mobile 无缝访问 NNAPI。此初始版本包括对核心功能和操作符集的完整功能支持,Google 和 Facebook 将在未来几个月内努力扩展功能。

链接

PyTorch Mobile GPU 支持

在 GPU 上进行推理可以在许多模型类型上提供卓越的性能,特别是那些利用高精度浮点运算的模型。利用 GPU 执行 ML 模型(如 Qualcomm、Mediatek 和 Apple SOC 中的 GPU)可以分担 CPU 负载,从而解放移动 CPU 用于非 ML 用例。此初始原型级别的设备上 GPU 支持是通过 iOS 的 Metal API 规范和 Android 的 Vulkan API 规范实现的。由于此特性处于早期阶段:性能尚未优化,模型覆盖范围有限。我们预计这种情况将在 2021 年显著改善,并且希望听取您关于希望在哪些模型和设备上看到性能改进的意见。

链接

Linux ARM64 构建

我们现在将为 Linux 上的 ARM64 设备提供原型级别的 PyTorch 构建。我们看到社区中 ARM 的使用越来越多,例如 Raspberry Pi 和 Graviton(2) 实例,它们分别涵盖边缘和服务器。此特性可通过我们的每夜构建版本获取。

我们重视您对这些特性的反馈,并期待与您合作,不断进一步改进它们!

谢谢,

PyTorch 团队