AI Edge Torch - PyTorch 模型转换为 TensorFlow Lite 模型 (.tflite)

google-ai-edge
https://github.com/google-ai-edge

1. AI Edge Torch

ai-edge-torch
https://github.com/google-ai-edge/ai-edge-torch

AI Edge Torch is a python library that supports converting PyTorch models into a .tflite format, which can then be run with TensorFlow Lite and MediaPipe. This enables applications for Android, iOS and IOT that can run models completely on-device.

AI Edge Torch offers broad CPU coverage, with initial GPU and NPU support. AI Edge Torch seeks to closely integrate with PyTorch, building on top of torch.export() and providing good coverage of Core ATen operators.

AI Edge Torch 可以将 PyTorch 编写的模型直接转换成 TFLite 格式 (.tflite)，且有着优异的模型覆盖率和 CPU 性能。

1.1. PyTorch Converter

Here are the steps needed to convert a PyTorch model to a TFLite flatbuffer:

import torch
import torchvision
import ai_edge_torch

# Use resnet18 with pre-trained weights.
resnet18 = torchvision.models.resnet18(torchvision.models.ResNet18_Weights.IMAGENET1K_V1)
sample_inputs = (torch.randn(1, 3, 224, 224),)

# Convert and serialize PyTorch model to a tflite flatbuffer. Note that we
# are setting the model to evaluation mode prior to conversion.
edge_model = ai_edge_torch.convert(resnet18.eval(), sample_inputs)
edge_model.export("resnet18.tflite")

docs/pytorch_converter/getting_started.ipynb
https://github.com/google-ai-edge/ai-edge-torch/blob/main/docs/pytorch_converter/getting_started.ipynb

docs/pytorch_converter/README.md
https://github.com/google-ai-edge/ai-edge-torch/blob/main/docs/pytorch_converter/README.md

1.2. Generative API

The AI Edge Torch Generative API is a Torch native library for authoring mobile-optimized PyTorch Transformer models, which can be converted to TFLite, allowing users to easily deploy Large Language Models (LLMs) on mobile devices. Users can convert the models using the AI Edge Torch PyTorch Converter, and run them via the TensorFlow Lite runtime.

AI Edge Torch Generative API 可以将开发者用 PyTorch 编写的高性能大语言模型 (LLM) 部署至 TensorFlow Lite (TFLite) 运行时，从而无缝地将新的设备端生成式 AI 模型部署到边缘设备上。

Example transformer models (decoder-only LLMs)
https://github.com/google-ai-edge/ai-edge-torch/tree/main/ai_edge_torch/generative/examples

Mobile app developers can also use the Edge Generative API to integrate PyTorch LLMs directly with the MediaPipe LLM Inference API for easy integration within their application code.

LLM Inference guide - LLM 推断指南
https://ai.google.dev/edge/mediapipe/solutions/genai/llm_inference

AI Edge Torch Generative API
https://github.com/google-ai-edge/ai-edge-torch/tree/main/ai_edge_torch/generative

2. Installation

2.1. Update LD_LIBRARY_PATH if necessary

Torch XLA builds a shared library, _XLAC.so that needs to link to the version of Python it was built with (currently 3.10 or 3.11). In order to ensure that import _XLAC can succeed, update the LD_LIBRARY_PATH to the lib directory of your Python environment:

export LD_LIBRARY_PATH=<path to Python installation>/lib:$LD_LIBRARY_PATH

References

[1] Yongqiang Cheng, https://yongqiang.blog.csdn.net/
[2] AI Edge Torch, https://github.com/google-ai-edge/ai-edge-torch