Converter modelos PyTorch para o TF Lite

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