AI Edge Torch es una biblioteca que te permite convertir modelos de PyTorch en un formato de .tflite
, lo que te permite ejecutar esos modelos con TensorFlow Lite y MediaPipe.
Esto es especialmente útil para los desarrolladores que crean apps para dispositivos móviles que ejecutan modelos completamente en el dispositivo. AI Edge Torch ofrece una amplia cobertura de CPU, con compatibilidad inicial con GPU y NPU.
Para comenzar a convertir modelos de PyTorch a TF Lite, usa la guía de inicio rápido del conversor de PyTorch. Para obtener más información, consulta el repositorio de GitHub de AI Edge Torch.
Si quieres convertir, de manera específica, modelos grandes de lenguaje (LLM) o modelos basados en transformadores, usa la API de Generative Torch, que controla los detalles de conversión específicos de transformadores, como la creación y la cuantización de modelos.
Flujo de trabajo de las conversiones
En los siguientes pasos, se muestra una conversión simple de extremo a extremo de un modelo de PyTorch en TensorFlow Lite.
Importa AI Edge Torch
Comienza por importar el paquete de pip de AI Edge Torch (ai-edge-torch
) junto con PyTorch.
import ai_edge_torch
import torch
Para este ejemplo, también necesitamos los siguientes paquetes:
import numpy
import torchvision
Inicializa y convierte el modelo
Convertiremos ResNet18, un modelo popular de reconocimiento de imágenes.
resnet18 = torchvision.models.resnet18(torchvision.models.ResNet18_Weights.IMAGENET1K_V1).eval()
Usa el método convert
de la biblioteca de AI Edge Torch para convertir el modelo de PyTorch.
sample_input = (torch.randn(1, 3, 224, 224),)
edge_model = ai_edge_torch.convert(resnet18.eval(), sample_input)
Usa el modelo
Después de convertir el modelo de Pytorch, puedes ejecutar inferencias con el nuevo modelo de TF Lite convertido.
output = edge_model(*sample_inputs)
Puedes exportar y guardar el modelo convertido en formato .tflite
para usarlo más adelante.
edge_model.export('resnet.tflite')