המרת מודלים של PyTorch ל-LiteRT

AI Edge Torch היא ספרייה שמאפשרת להמיר מודלים של PyTorch ל.tflite כך שתוכלו להריץ את המודלים האלה באמצעות LiteRT ו-MediaPipe. האפשרות הזו שימושית במיוחד למפתחים שיוצרים אפליקציות לנייד שמפעילות מודלים לגמרי במכשיר. עם AI Edge Torch נהנים מכיסוי רחב יותר של המעבד (CPU), עם GPU ראשוני ותמיכה ב-NPU.

כדי להתחיל בהמרת מודלים של PyTorch ל-LiteRT, יש להשתמש בממיר Pytorch במדריך למתחילים. מידע נוסף זמין ב-AI Edge Torch מאגר ב-GitHub

אם אתם ממירים באופן ספציפי מודלים גדולים של שפה (LLMs) או שמבוססים על טרנספורמרים, משתמשים ב-Generative Torch API, המטפל בפרטי המרות ספציפיים לטרנספורמר, כמו יצירת מודלים יצירת קונטיינרים.

תהליך עבודה של המרות

השלבים הבאים מדגימים המרה פשוטה מקצה לקצה של PyTorch ל-LiteRT.

ייבוא של AI Edge לפיד

קודם כול צריך לייבא את חבילת ה-PIP של AI Edge (ai-edge-torch), יחד עם פיטור.

import ai_edge_torch
import torch

לצורך הדוגמה הזו, צריך גם את החבילות הבאות:

import numpy
import torchvision

אתחול והמרה של המודל

אנחנו נמיר ResNet18, מודל פופולרי לזיהוי תמונות.

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

אפשר להשתמש בשיטה convert מספריית AI Edge Torch כדי להמיר את PyTorch מודל טרנספורמר.

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

שימוש במודל

אחרי שממירים את מודל Pytorch, אפשר להריץ מסקנות את מודל LiteRT שהומר.

output = edge_model(*sample_inputs)

ניתן לייצא ולשמור את המודל שהומר בפורמט .tflite לשימוש עתידי לשימוש.

edge_model.export('resnet.tflite')