AI Edge Torch: An Overview
AI Edge Torch is a Python library developed to transform PyTorch machine learning models into TensorFlow Lite files, allowing these models to be operational on platforms such as Android, iOS, and IoT devices. This transformation enables applications to execute models directly on devices, enhancing performance and broadening their accessibility. The library primarily supports CPU execution, with initial capabilities for GPU and NPU support. AI Edge Torch aims to effectively integrate with PyTorch, utilizing its torch.export() function and offering extensive support for Core ATen operators.
PyTorch Converter
The PyTorch Converter is a key component of AI Edge Torch, allowing users to convert PyTorch models into TFLite flatbuffers. Here's a basic example of how one might convert a model:
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.
edge_model = ai_edge_torch.convert(resnet18.eval(), sample_inputs)
edge_model.export("resnet18.tflite")
This process involves setting the model to evaluation mode before conversion, and the command exports the model into a .tflite file format. For beginners, a Jupyter notebook is available for an interactive experience, which can also be run using Google Colab. More technical specifics are provided in the documentation.
Generative API
The Generative API within AI Edge Torch focuses on enabling the creation of transformer models optimized for mobile devices using PyTorch. These models can be transformed into TFLite files, similar to other models, thus facilitating the deployment of Large Language Models (LLMs) on mobile platforms. Developers are supported in using PyTorch models with the MediaPipe LLM Inference API for seamless application integration. This functionality can currently be employed with CPU, but support for GPU and NPU is being planned.
For further technical insights, users can explore additional examples here or refer to more documentation here.
Current Build Status
The project’s build status indicates ongoing development and maintenance:
- Generative API (Linux):
- Model Coverage (Linux):
- Unit tests (Linux):
- Nightly Release:
Installation Instructions
System Requirements
- Python: Versions 3.9, 3.10, 3.11
- Operating System: Linux
- Frameworks: PyTorch (≥ 2.4.0), TensorFlow (latest nightly build)
Installing in a Virtual Environment
Users should create a Python virtual environment for installation:
python -m venv --prompt ai-edge-torch venv
source venv/bin/activate
To install the latest stable release, run:
pip install ai-edge-torch
For the nightly version, use:
pip install ai-edge-torch-nightly
Library Path Considerations
Torch XLA requires linking to the Python version it was built with. Update the LD_LIBRARY_PATH to point to the correct Python library path:
export LD_LIBRARY_PATH=<path to Python installation>/lib:$LD_LIBRARY_PATH
For complete release information, visit GitHub releases or check PyPi history.
Contributing and Support
For those interested in contributing, detailed guidelines are available in the contribution documentation. Should users encounter any issues or have queries, they are encouraged to create a GitHub issue.
AI Edge Torch provides robust tools for on-device machine learning model deployment, enhancing both mobile and IoT applications by leveraging the strengths of PyTorch and TensorFlow Lite.