AI Edge Torch: 将PyTorch模型部署到移动设备上的高性能推理框架

AI Edge Torch是Google最近发布的一个开源Python库,旨在简化将PyTorch模型部署到移动设备和IoT设备上的过程。作为Google AI Edge生态系统的一部分,这个库为开发者提供了一种便捷的方式,可以将PyTorch模型转换为TensorFlow Lite (TFLite)格式,并在移动端和IoT设备上实现高性能的本地推理。

主要特性和优势

AI Edge Torch的主要特性和优势包括:

直接的PyTorch集成 - 提供与PyTorch原生感觉一致的API,使转换过程变得简单直观。
优秀的CPU性能和初步的GPU支持 - 在CPU上提供出色的推理性能,同时也开始支持GPU加速。
广泛的模型验证 - 已在70多个来自torchvision、timm、torchaudio和HuggingFace的模型上进行了验证测试。
高覆盖率 - 支持超过70%的PyTorch core_aten算子。
与现有TFLite运行时兼容 - 无需更改部署代码即可使用。
支持Model Explorer可视化 - 在工作流程的多个阶段提供模型可视化功能。

PyTorch转换器

AI Edge Torch的核心功能是PyTorch转换器,它可以将PyTorch模型转换为TFLite格式。使用这个转换器非常简单,只需几行代码即可完成:

import torch
import torchvision
import ai_edge_torch

# 使用预训练的ResNet18模型
resnet18 = torchvision.models.resnet18(torchvision.models.ResNet18_Weights.IMAGENET1K_V1)
sample_inputs = (torch.randn(1, 3, 224, 224),)

# 将PyTorch模型转换为TFLite格式
edge_model = ai_edge_torch.convert(resnet18.eval(), sample_inputs)
edge_model.export("resnet18.tflite")

这个简单的示例展示了如何将预训练的ResNet18模型转换为TFLite格式。开发者可以轻松地将这个过程应用到自己的自定义PyTorch模型上。

Generative API

除了基本的PyTorch转换器之外,AI Edge Torch还提供了一个Generative API,专门用于处理大型语言模型(LLMs)和基于Transformer的模型。这个API支持模型创作和量化,以实现更好的设备端性能。

Generative API的主要特点包括:

支持自定义PyTorch Transformer模型的创作
提供模型量化功能,优化设备端性能
与MediaPipe LLM推理API集成,便于在应用程序中使用

目前,Generative API仅支持CPU,但计划在未来添加GPU和NPU支持。

AI Edge Torch Architecture

安装和使用

要开始使用AI Edge Torch,你需要以下环境:

Python 3.9、3.10或3.11
Linux操作系统
PyTorch 2.4.0或更高版本
TensorFlow nightly版本

你可以通过pip安装AI Edge Torch:

pip install ai-edge-torch

如果你想尝试最新的开发版本,可以安装nightly版本:

pip install ai-edge-torch-nightly

实际应用案例

AI Edge Torch已经在一些知名公司的实际项目中得到应用。例如,Shopify正在使用AI Edge Torch来实现产品图像的本地背景去除功能,这将在Shopify应用的即将发布的版本中推出。

Shopify的工程总监Mustapha Ali表示:"将PyTorch模型转换为在Android上本地运行曾经是一个复杂的过程。Google的新工具简化了这一过程,使我们能够快速创建适用于移动设备的PyTorch模型。"

硬件加速与合作伙伴

Google还与多家硬件合作伙伴合作,为AI Edge Torch提供硬件加速支持。这些合作伙伴包括Arm、Google Tensor G3、MediaTek、Qualcomm和Samsung System LSI。通过这些合作,AI Edge Torch在不同的硬件平台上都能获得优化的性能。

特别值得一提的是,Qualcomm推出了新的TensorFlow Lite delegate,现已开放供所有开发者使用。这个delegate支持AI Edge Torch测试集中的大多数模型,并可以在Qualcomm的DSP和神经处理单元上实现显著的性能提升,相比CPU平均可提速20倍,相比GPU可提速5倍。