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

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

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

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

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

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

ייבוא לפיד AI Edge

כדי להתחיל, צריך לייבא את חבילת ה-Pyp של 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()

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

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

שימוש במודל

אחרי ההמרה של מודל Pytorch, אפשר להריץ הסקת מסקנות במודל TF Lite החדש שעבר המרה.

output = edge_model(*sample_inputs)

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

edge_model.export('resnet.tflite')