AI Edge Torch is a Python library designed to convert PyTorch models into the TFLite format, enabling on-device execution for Android, iOS, and IoT applications. It targets developers building mobile and edge AI solutions, offering broad CPU support with initial GPU and NPU capabilities, and a specialized Generative API for LLMs.

How It Works

The library leverages torch.export() to convert PyTorch models, aiming for broad coverage of Core ATen operators. This approach integrates closely with the PyTorch ecosystem. For LLMs, the Generative API facilitates model authoring and quantization for optimized on-device performance, with conversion to TFLite for execution via the TensorFlow Lite runtime or direct integration with the MediaPipe LLM Inference API.

Quick Start & Requirements

• Install: pip install ai-edge-torch or pip install ai-edge-torch-nightly
• Requirements: Python >= 3.10, Linux, PyTorch, TensorFlow.
• Documentation: PyTorch converter, Generative API, Jupyter Notebook

Highlighted Details

• Supports conversion of PyTorch models to TFLite format for on-device inference.
• Generative API specifically targets LLMs and transformer models for mobile deployment.
• Offers broad CPU coverage with planned GPU and NPU support.
• Integrates with MediaPipe LLM Inference API for mobile application development.

Maintenance & Community

• Users can seek help by creating a GitHub issue.
• Contribution documentation is available.

Licensing & Compatibility

• License details are not explicitly stated in the README.

Limitations & Caveats

The PyTorch converter is in Beta, and the Generative API is in Alpha. The Generative API is currently CPU-only, with GPU and NPU support planned.