AI Edge Torch

AI Edge Torch 是一个 Python 库，支持将 PyTorch 模型转换为 .tflite 格式，然后可以使用 TensorFlow Lite 和 MediaPipe 运行。这使得 Android、iOS 和 IoT 应用程序能够完全在设备上运行模型。AI Edge Torch 提供广泛的 CPU 支持，并初步支持 GPU 和 NPU。AI Edge Torch 致力于与 PyTorch 紧密集成，基于 torch.export() 构建，并提供对核心 ATen 运算符的良好覆盖。

要开始将 PyTorch 模型转换为 TF Lite，请参阅 PyTorch 转换器 部分的更多详细信息。对于大型语言模型（LLM）和基于 Transformer 的模型，生成式 API 支持模型创作和量化，以提高设备上的性能。

虽然同属一个 PyPi 包，但 PyTorch 转换器是 Beta 版本，而生成式 API 是 Alpha 版本。请查看发布说明以获取更多信息。

PyTorch 转换器

以下是将 PyTorch 模型转换为 TFLite flatbuffer 的步骤：

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 flatbuffer。注意在转换前
# 我们将模型设置为评估模式。
edge_model = ai_edge_torch.convert(resnet18.eval(), sample_inputs)
edge_model.export("resnet18.tflite")

入门指南 Jupyter notebook 提供了转换过程的初步演示，可以在 Google Colab 中尝试。

PyTorch 转换器的更多技术细节请见此处。

生成式 API

AI Edge Torch 生成式 API 是一个用于创作移动优化 PyTorch Transformer 模型的 Torch 原生库，这些模型可以转换为 TFLite，让用户轻松在移动设备上部署大型语言模型（LLM）。用户可以使用 AI Edge Torch PyTorch 转换器转换模型，并通过 TensorFlow Lite 运行时运行它们。参见此处。

移动应用开发者还可以使用 Edge 生成式 API 将 PyTorch LLM 直接与 MediaPipe LLM 推理 API 集成，以便在应用程序代码中轻松整合。参见此处。

更详细的文档可以在这里找到。

生成式 API 目前仅支持 CPU，计划未来支持 GPU 和 NPU。未来的另一个方向是与 PyTorch 社区合作，确保常用的 transformer 抽象可以直接支持，无需重新创作。

构建状态

构建类型	状态
生成式 API (Linux)
模型覆盖率 (Linux)
单元测试 (Linux)
每日构建

安装

要求和依赖

• Python 版本：3.9、3.10、3.11
• 操作系统：Linux
• PyTorch:
• TensorFlow:

Python 虚拟环境

设置 Python 虚拟环境：

python -m venv --prompt ai-edge-torch venv
source venv/bin/activate

可以通过以下命令安装最新的稳定版本：

pip install ai-edge-torch

或者，可以通过以下命令安装每日构建版本：

pip install ai-edge-torch-nightly

如有必要，更新 LD_LIBRARY_PATH

Torch XLA 构建了一个共享库 _XLAC.so，需要链接到其构建时使用的 Python 版本（目前是 3.10 或 3.11）。为确保 import _XLAC 能够成功，请将 LD_LIBRARY_PATH 更新到您的 Python 环境的 lib 目录：

export LD_LIBRARY_PATH=<Python 安装路径>/lib:$LD_LIBRARY_PATH

• 版本发布列表可以在这里查看。
• 完整的 PyPi 发布列表（包括每日构建）可以在这里查看。

贡献

请参阅我们的贡献文档。

获取帮助

如有任何问题，请创建 GitHub issue。