Converter modelos PyTorch para o LiteRT

A AI Edge Torch é uma biblioteca que permite converter modelos PyTorch em um .tflite , o que permite a execução desses modelos com LiteRT e MediaPipe. Isso é especialmente útil para desenvolvedores que criam apps para dispositivos móveis que executam modelos. totalmente no dispositivo. O AI Edge Torch oferece ampla cobertura de CPU, com GPU inicial e NPU.

Para começar a converter modelos PyTorch para LiteRT, use o conversor PyTorch guia de início rápido. Para mais informações, consulte o AI Edge Torch Repositório do GitHub.

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 de modelo e e quantização.

Fluxo de trabalho de conversão

As etapas a seguir demonstram uma conversão simples de ponta a ponta de um PyTorch para LiteRT.

Importar AI Edge Torch

Comece importando o pacote pip AI Edge Torch (ai-edge-torch) com 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 (link em inglês) um modelo conhecido de reconhecimento de imagens.

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

Usar o método convert da biblioteca AI Edge Torch para converter o PyTorch um modelo de machine learning.

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 LiteRT convertido.

output = edge_model(*sample_inputs)

É possível exportar e salvar o modelo convertido no formato .tflite para uso futuro usar.

edge_model.export('resnet.tflite')