PyTorch-Modelle in LiteRT konvertieren

AI Edge Torch ist eine Bibliothek, mit der Sie PyTorch-Modelle in ein .tflite-Objekt konvertieren können sodass Sie diese Modelle mit LiteRT und MediaPipe ausführen können. Dies ist besonders hilfreich für Entwickler von mobilen Apps, die Modelle ausführen vollständig auf dem Gerät. AI Edge Torch bietet eine große CPU-Abdeckung mit anfänglicher GPU und NPU unterstützt.

Verwenden Sie zum Konvertieren von PyTorch-Modellen in LiteRT den PyTorch-Konverter. Kurzanleitung. Weitere Informationen findest du im Artikel zu AI Edge Torch GitHub-Repository.

Wenn Sie speziell Large Language Models (LLMs) konvertieren oder Transformer-basierten Modellen die Generative Torch API, das transformerspezifische Konvertierungsdetails wie Modellerstellung und Quantisierung.

Conversion-Workflow

Die folgenden Schritte zeigen eine einfache End-to-End-Konvertierung eines PyTorch-Moduls. in LiteRT umwandeln.

AI Edge Torch importieren

Importieren Sie zuerst das pip-Paket von AI Edge Torch (ai-edge-torch) zusammen mit PyTorch.

import ai_edge_torch
import torch

Für dieses Beispiel sind außerdem die folgenden Pakete erforderlich:

import numpy
import torchvision

Modell initialisieren und konvertieren

Wir konvertieren ResNet18 ein beliebtes Bilderkennungsmodell.

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

Verwenden Sie die Methode convert aus der AI Edge Torch-Bibliothek, um das PyTorch-Objekt zu konvertieren Modell.

sample_input = (torch.randn(1, 3, 224, 224),)
edge_model = ai_edge_torch.convert(resnet18.eval(), sample_input)

Modell verwenden

Nach der Konvertierung des Pytorch-Modells können Sie mit dem neuen konvertierten LiteRT-Modells.

output = edge_model(*sample_inputs)

Sie können das konvertierte Modell für später exportieren und im .tflite-Format speichern verwenden.

edge_model.export('resnet.tflite')