AI Edge Torch to biblioteka, która umożliwia konwertowanie modeli PyTorch do formatu .tflite, co pozwala na ich uruchamianie za pomocą LiteRT i MediaPipe.
Jest to szczególnie przydatne dla deweloperów tworzących aplikacje mobilne, które uruchamiają modele całkowicie na urządzeniu. AI Edge Torch zapewnia szerokie pokrycie procesora z pierwotnym wsparciem GPU i NPU.
Aby rozpocząć konwertowanie modeli PyTorch na LiteRT, skorzystaj z Pytorch converter quickstart. Więcej informacji znajdziesz w repozytorium GitHub Torch AI Edge.
Jeśli konwertujesz duże modele językowe (LLM) lub modele oparte na transformatorach, użyj interfejsu Generative Torch API, który obsługuje szczegóły konwersji związane z transformerami, takie jak tworzenie modelu i kwantyzacja.
Proces konwersji
Te kroki pokazują prostą konwersję modelu PyTorch na LiteRT.
Importowanie AI Edge Torch
Zacznij od zaimportowania pakietu pip AI Edge Torch (ai-edge-torch) wraz z PyTorch.
import ai_edge_torch
import torch
W tym przykładzie potrzebujemy też tych pakietów:
import numpy
import torchvision
Inicjowanie i konwertowanie modelu
Przekształcimy ResNet18, popularny model rozpoznawania obrazu.
resnet18 = torchvision.models.resnet18(torchvision.models.ResNet18_Weights.IMAGENET1K_V1).eval()
Aby przekonwertować model PyTorch, użyj metody convert z biblioteki AI Edge Torch.
sample_input = (torch.randn(1, 3, 224, 224),)
edge_model = ai_edge_torch.convert(resnet18.eval(), sample_input)
Korzystanie z modelu
Po przekonwertowaniu modelu Pytorch możesz przeprowadzać wnioskowanie za pomocą nowego przekonwertowanego modelu LiteRT.
output = edge_model(*sample_inputs)
Konwertowany model możesz wyeksportować i zapisz w formacie .tflite, aby korzystać z niego w przyszłości.
edge_model.export('resnet.tflite')