Converti i modelli PyTorch in TF Lite

AI Edge Torch è una libreria che ti consente di convertire i modelli PyTorch in un formato .tflite, per eseguirli con TensorFlow Lite e MediaPipe. Ciò è particolarmente utile per gli sviluppatori che creano app mobile che eseguono modelli completamente sul dispositivo. AI Edge Torch offre un'ampia copertura della CPU, con supporto iniziale di GPU e NPU.

Per iniziare a convertire i modelli PyTorch in TF Lite, consulta la guida rapida al convertitore di PyTorch. Per maggiori informazioni, consulta il repository GitHub di AI Edge Torch.

Se stai convertendo in modo specifico modelli linguistici di grandi dimensioni (LLM) o modelli basati su trasformatore, utilizza l'API Generative Torch, che gestisce i dettagli di conversione specifici del trasformatore, come la creazione e la quantizzazione dei modelli.

Flusso di lavoro della conversione

I passaggi seguenti mostrano una semplice conversione end-to-end di un modello PyTorch in TensorFlow Lite.

Importa AI Edge Torch

Inizia importando il pacchetto pip AI Edge Torch (ai-edge-torch), insieme a PyTorch.

import ai_edge_torch
import torch

Per questo esempio, sono richiesti anche i seguenti pacchetti:

import numpy
import torchvision

Inizializza e converti il modello

Convertiremo ResNet18, un popolare modello di riconoscimento delle immagini.

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

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

Utilizza il modello

Dopo aver convertito il modello Pytorch, puoi eseguire inferenze con il nuovo modello TF Lite 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')