Convierte modelos de PyTorch a TF Lite

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')