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

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

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

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

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

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

ייבוא AI Edge Torch

מתחילים בייבוא חבילת ה-pip של AI Edge Torch‏ (ai-edge-torch) יחד עם PyTorch.

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