Converter modelos PyTorch para o LiteRT

O AI Edge Torch é uma biblioteca que permite converter modelos do PyTorch em um formato .tflite, para que você possa executá-los com o LiteRT e o MediaPipe. Isso é especialmente útil para desenvolvedores que criam apps para dispositivos móveis que executam modelos completamente no dispositivo. A AI Edge Torch oferece ampla cobertura de CPU, com suporte inicial de GPU e NPU.

Para começar a converter modelos do PyTorch para o LiteRT, use o Pytorch converter quickstart. Para mais informações, consulte o repositório do GitHub da AI Edge Torch.

Se você estiver convertendo especificamente modelos de linguagem grandes (LLMs) ou modelos baseados em transformador, use a API Generative Torch, que processa detalhes de conversão específicos do transformador, 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 LiteRT.

Importar IA Edge Torch

Comece importando o pacote pip AI Edge Torch (ai-edge-torch) com o PyTorch.

import ai_edge_torch
import torch

Para este exemplo, também precisamos dos seguintes pacotes:

import numpy
import torchvision

Inicializar e converter o modelo

Vamos converter ResNet18, um modelo de reconhecimento de imagem conhecido.

resnet18 = torchvision.models.resnet18(torchvision.models.ResNet18_Weights.IMAGENET1K_V1).eval()

Use o método convert da biblioteca 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, você pode executar inferências com o novo modelo LiteRT 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')