Convierte modelos de PyTorch en LiteRT

AI Edge Torch es una biblioteca que te permite convertir modelos de PyTorch a un formato .tflite, lo que te permite ejecutar esos modelos con LiteRT 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 LiteRT, usa la Guía de inicio rápido del convertidor de PyTorch. Para obtener más información, consulta el repositorio de GitHub de AI Edge Torch.

Si conviertes específicamente modelos de lenguaje extenso (LLM) o basados en transformadores, usa la API de Generative Torch, que controla los detalles de conversión específicos del transformador, como la autoría y la quantización del modelo.

Flujo de trabajo de conversión

En los siguientes pasos, se muestra una conversión de extremo a extremo simple de un modelo de PyTorch a LiteRT.

Importa AI Edge Torch

Comienza por importar el paquete de pip AI Edge Torch (ai-edge-torch) junto con PyTorch.

import ai_edge_torch
import torch

Para este ejemplo, también requerimos 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 LiteRT convertido.

output = edge_model(*sample_inputs)

Puedes exportar y guardar el modelo convertido en formato .tflite para usarlo en el futuro.

edge_model.export('resnet.tflite')