AI Edge Torch：为 Pytorch 模型在移动设备上提供高性能推理

五月 14, 2024

Cormac Brick Principal Engineer

Advait Jain Software Engineer

Haoliang Zhang Software Engineer

我们很高兴地宣布 Google AI Edge Torch - 一个从 PyTorch 到 TensorFlow Lite (TFLite) 运行时的直接路径，它具有出色的模型覆盖率和 CPU 性能。 TFLite 已经可以与用 Jax、Keras 和 TensorFlow 编写的模型配合使用，我们现在还添加了 PyTorch。正如我们承诺过的，我们会让用户拥有更广泛的选择。

这项新产品现已作为 Google AI Edge 的一部分提供n。这是一套工具，可以轻松访问现成的机器学习（ML）任务、框架，使您能够构建 ML 管线并运行流行的大语言模型（LLM）和自定义模型。这些全部能在设备上完成。这是涵盖 Google AI Edge 版本的一系列博客文章中的第一篇，这些文章将帮助开发人员构建支持 AI 的功能，并轻松将其部署在多个平台上。

AI Edge Torch 现已发布测试版，其中包括：

直接的 PyTorch 集成
出色的 CPU 性能和初始 GPU 支持
已在 torchvision、timm、torchaudio 和 HuggingFace 等70多个模型上进行了验证
支持 PyTorch 中 > 70% 的 core_aten 运算符
与现有 TFLite 运行时兼容，无需更改部署代码
支持工作流程中多个阶段的 Model Explorer 可视化。

简单的、以 PyTorch 为中心的体验

全新构建的 Google AI Edge Torch，旨在为 PyTorch 社区提供出色的体验，给用户提供原生感的 API 及提供简便的转换路径。

import torchvision
import ai_edge_torch
 
# Initialize model
resnet18 = torchvision.models.resnet18().eval()
 
# Convert
sample_input = (torch.randn(4, 3, 224, 224),)
edge_model = ai_edge_torch.convert(resnet18, sample_input)
 
# Inference in Python
output = edge_model(*sample_input)
 
# Export to a TfLite model for on-device deployment
edge_model.export('resnet.tflite'))

在内部实现中，ai_edge_torch.convert() 集成了使用 torch.export 的 TorchDynamo - 这是 PyTorch 2.x 用于将 PyTorch 模型导出到不同环境中运行的标准化模型的表示形式。我们当前的实现支持超过 60% 的 core_aten 运算符，我们计划在构建 ai_edge_torch 1.0 版本时大幅增加这一比例。我们提供了展示 PT2E 量化（PyTorch2 原生的量化方法）的示例，以实现轻松的量化工作流程。我们很高兴聆听到 PyTorch 社区的声音，了解到社区希望能将 PyTorch 引入到广泛的设备当中来改善开发者的体验。

覆盖率 & 性能

在此版本之前，许多开发者使用社区提供的路径（例如 ONNX2TF）在 TFLite 上启用 PyTorch 模型。我们开发 AI Edge Torch 的目标是减少开发人员的摩擦，提供出色的模型覆盖率，并继续履行我们在 Android 设备上提供一流性能的使命。

在覆盖范围方面，我们的测试显示，与现有工作流程相比，我们的流程在给定模型集上有着显着改进，尤其是 ONNX2TF

Table showing performance improvement in existing workflows over defined set models

在性能方面，我们的测试显示出与 ONNX2TF 基准一致的性能，同时还显示出比 ONNX 运行时明显更好的性能：

Table showing performance with ONNX2TF baseline

图片显示了 ONNX 涵盖的模型子集上每个模型的详细性能：

Chart showing per model TFLite latency relative to ONNX

图：与 ONNX 相比每个网络的推理延迟，测量设备为 Pixel8，精度为 fp32，为提高可重复性 XNNPACK 固定为4个线程，20 次预热迭代后运行100次的平均值

早期采购合作伙伴

在过去的几个月中，我们与 Shopify、Adobe 和 Niantic 等早期采购商密切合作，以改进我们的 PyTorch 支持。 Shopify 团队已使用 ai_edge_torch 来执行设备上产品图像背景的删除，并将运用在即将发布的 Shopify 应用中。

Quote image with text reads "Converting PyTorch models to run locally on Android was complex. Google's new tools simplify this, enabling fast creation of mobile-ready PyTorch models - Mustapha Ali, Shopify, Director of Engineering

硬件合作伙伴&代理

我们还与合作伙伴密切合作，致力于跨 CPU、GPU 和加速器的硬件支持，其中包括 Arm、Google Tensor G3、MediaTek、Qualcomm 和 Samsung System LSI。通过这些合作伙伴，我们提高了性能和覆盖范围，并在加速器委托上验证了 PyTorch 生成的 TFLite 文件。

我们还很高兴与大家共同推出 Qualcomm（高通）的全新 TensorFlow Lite 代理。该代理现已在此公开提供给任何开发者使用。 TFLite Delegates 是附加软件模块，有助于加速 GPU 和硬件加速器上的执行。这个新的 QNN delegate 支持 PyTorch Beta 测试集中的大多数模型，同时为广泛的 Qualcomm 芯片提供支持，并通过利用 Qualcomm 的 DSP 和神经处理单元，相对于 CPU (20x) 和 GPU(5x) 提供显着的平均加速。为了方便测试，Qualcomm 最近还发布了新的人工智能中心（AI Hub）。 Qualcomm AI Hub 是一项云服务，使开发人员能够针对广泛的 Android 设备池测试 TFLite 模型，并使用 QNN 委托提供在不同设备上都可见的性能增益。

未来计划

在接下来的几个月中，我们将继续开源我们的进度，随着我们构建 1.0 版本，发布版本将扩大模型覆盖范围、改进 GPU 支持并启用新的量化模式。在本系列的第 2 部分中，我们将深入了解 AI Edge Torch Generative API，它使开发人员能够将自定义的 GenAI（生成式人工智能）模型引入边缘设备并具有出色的性能。

我们感谢所有早期试用客户提供的宝贵意见反馈。这些反馈帮助我们及早发现错误并确保提供给开发者流畅的使用体验。我们还要感谢 XNNPACK 的硬件合作伙伴和生态系统贡献者，他们帮助我们提高了各种设备的性能。我们还要感谢来自 Pytorch 社区的广泛的指导和支持。

致谢

^{我们要感谢所有为这项工作做出贡献的团队成员：Aaron Karp, Advait Jain, Akshat Sharma, Alan Kelly, Arian Arfaian, Chun-nien Chan, Chuo-Ling Chang, Claudio Basille, Cormac Brick, Dwarak Rajagopal, Eric Yang, Gunhyun Park, Han Qi, Haoliang Zhang, Jing Jin, Juhyun Lee, Jun Jiang, Kevin Gleason, Khanh LeViet, Kris Tonthat, Kristen Wright, Lu Wang, Luke Boyer, Majid Dadashi, Maria Lyubimtseva, Mark Sherwood, Matthew Soulanille, Matthias Grundmann, Meghna Johar, Milad Mohammadi, Na Li, Paul Ruiz, Pauline Sho, Ping Yu, Pulkit Bhuwalka, Ram Iyengar, Sachin Kotwani, Sandeep Dasgupta, Sharbani Roy, Shauheen Zahirazami, Siyuan Liu, Vamsi Manchala, Vitalii Dziuba, Weiyi Wang, Wonjoo Lee, Yishuang Pang, Zoe Wang, 以及 StableHLO 团队.}