O AI Edge Torch é uma biblioteca que permite converter modelos PyTorch em um formato .tflite
, permitindo que você execute esses modelos com o TensorFlow Lite e o MediaPipe.
Isso é especialmente útil para desenvolvedores que criam apps para dispositivos móveis que executam modelos
completamente no dispositivo. O AI Edge Torch oferece ampla cobertura de CPU, com suporte inicial a GPU e NPU.
Para começar a converter modelos PyTorch para o TF Lite, use o guia de início rápido do conversor do PyTorch. Para mais informações, consulte o repositório do GitHub do AI Edge Torch.
Se você estiver convertendo especificamente modelos de linguagem grandes (LLMs) ou modelos baseados em transformadores, use a API Generative Torch, que processa detalhes de conversão específicos de transformadores, como criação e quantização de modelos.
Fluxo de trabalho de conversão
As etapas a seguir demonstram uma conversão completa e simples de um modelo PyTorch para o TensorFlow Lite.
Importar AI Edge Torch
Comece importando o pacote pip do AI Edge Torch (ai-edge-torch
) com o
PyTorch.
import ai_edge_torch
import torch
Para este exemplo, também são necessários os seguintes pacotes:
import numpy
import torchvision
Inicializar e converter o modelo
Converteremos o ResNet18, um modelo conhecido de reconhecimento de imagem.
resnet18 = torchvision.models.resnet18(torchvision.models.ResNet18_Weights.IMAGENET1K_V1).eval()
Use o método convert
da biblioteca do AI Edge Torch para converter o modelo PyTorch.
sample_input = (torch.randn(1, 3, 224, 224),)
edge_model = ai_edge_torch.convert(resnet18.eval(), sample_input)
use o modelo
Depois de converter o modelo Pytorch, é possível executar inferências com o novo modelo TF Lite convertido.
output = edge_model(*sample_inputs)
É possível exportar e salvar o modelo convertido no formato .tflite
para uso futuro.
edge_model.export('resnet.tflite')