AI Edge Torch è una libreria che consente di convertire i modelli PyTorch in un formato .tflite, in modo da poterli eseguire con LiteRT e MediaPipe.
Ciò è particolarmente utile per gli sviluppatori che creano app mobile che eseguono modelli completamente sul dispositivo. AI Edge Torch offre una copertura CPU ampia, con supporto iniziale per GPU
e NPU.
Per iniziare a convertire i modelli PyTorch in LiteRT, utilizza la guida introduttiva al convertitore PyTorch. Per ulteriori informazioni, consulta il repository GitHub di AI Edge Torch.
Se converti specificamente modelli linguistici di grandi dimensioni (LLM) o modelli basati su transformer, utilizza l'API Generative Torch, che gestisce i dettagli di conversione specifici dei transformer, come la creazione e la quantizzazione dei modelli.
Flusso di lavoro di conversione
I passaggi che seguono mostrano una semplice conversione end-to-end di un modello PyTorch in LiteRT.
Importa torcia AI Edge
Inizia importando il pacchetto pip AI Edge Torch (ai-edge-torch) insieme a
PyTorch.
import ai_edge_torch
import torch
Per questo esempio, sono necessari anche i seguenti pacchetti:
import numpy
import torchvision
Inizializza e converti il modello
Convertiremo ResNet18, un modello di riconoscimento delle immagini molto diffuso.
resnet18 = torchvision.models.resnet18(torchvision.models.ResNet18_Weights.IMAGENET1K_V1).eval()
Utilizza il metodo convert della libreria AI Edge Torch per convertire il
modello PyTorch.
sample_input = (torch.randn(1, 3, 224, 224),)
edge_model = ai_edge_torch.convert(resnet18.eval(), sample_input)
usa il modello
Dopo aver convertito il modello PyTorch, puoi eseguire inferenze con il nuovo modello LiteRT convertito.
output = edge_model(*sample_inputs)
Puoi esportare e salvare il modello convertito nel formato .tflite per un uso futuro.
edge_model.export('resnet.tflite')