תחילת העבודה עם LiteRT.js

זהו מדריך מקיף ל-LiteRT.js שמתאר את התהליך של המרת מודל PyTorch להרצה בדפדפן עם האצת WebGPU. בדוגמה הזו נעשה שימוש ב-ResNet18 למודל הראייה וב-TensorFlow.js לעיבוד מקדים ולעיבוד סופי.

במדריך הזה נסביר איך לבצע את השלבים הבאים:

  1. ממירים את מודל PyTorch ל-LiteRT באמצעות AI Edge Torch.
    1. מוסיפים את חבילת LiteRT לאפליקציית האינטרנט.
  2. טוענים את המודל.
  3. לכתוב לוגיקה לעיבוד מוקדם ולעיבוד אחרי.

המרת LiteRT

כדי להמיר מודל PyTorch לפורמט המתאים של .tflite, אפשר להשתמש במחברת PyTorch Converter. בקובץ ה-README של AI Edge Torch Converter יש מדריך מפורט על סוגי השגיאות שבהן אתם עשויים להיתקל ועל הדרכים לפתור אותן.

המודל צריך להיות תואם ל-torch.export.export, כלומר צריך להיות אפשר לייצא אותו באמצעות TorchDynamo. לכן, אסור שיהיו בו ענפים מותנים של Python שתלויים בערכי זמן הריצה בטנסורים. אם מופיעות השגיאות הבאות במהלך torch.export.export, אי אפשר לייצא את המודל באמצעות torch.export.export. בנוסף, במודל לא יכולים להיות מאפייני קלט או פלט דינמיים בטנסורים שלו. כולל מימד של קבוצת פריטים.

אפשר גם להתחיל עם מודל PyTorch שתואם ל-TensorRT או שאפשר לייצא אותו ל-ONNX:

  • גרסה של מודל שתואמת ל-TensorRT יכולה להיות נקודת התחלה טובה, כי חלק מהסוגים של המרות TensorRT דורשים גם שהמודלים יהיו ניתנים לייצוא ב-TorchDynamo. אם משתמשים בפעולות NVIDIA / CUDA במודל, צריך להחליף אותן בפעולות PyTorch רגילות.

  • מודל PyTorch שאפשר לייצא ל-ONNX יכול להיות נקודת התחלה טובה, אבל חלק ממודלי ONNX משתמשים ב-TorchScript במקום ב-TorchDynamo כדי לייצא. במקרה כזה, יכול להיות שלא ניתן לייצא את המודל באמצעות TorchDynamo (אבל סביר להניח שהוא קרוב יותר לזה מאשר קוד המודל המקורי).

מידע נוסף זמין במאמר המרת מודלים של PyTorch ל-LiteRT.

הוספת חבילת LiteRT

מתקינים את חבילת @litertjs/core מ-npm:

npm install @litertjs/core

מייבאים את החבילה וטוענים את קובצי ה-Wasm שלה:

import {loadLiteRt} from '@litertjs/core';

// They are located in node_modules/@litertjs/core/wasm/
// Serve them statically on your server.
await loadLiteRt(`your/path/to/wasm/`);

טעינת המודל

מייבאים ומפעילים את LiteRT.js ואת כלי ההמרה של LiteRT-TFJS. צריך גם לייבא את TensorFlow.js כדי להעביר טנסורים אל LiteRT.js.

import {CompileOptions, loadAndCompile, loadLiteRt, setWebGpuDevice} from '@litertjs/core';
import {runWithTfjsTensors} from '@litertjs/tfjs-interop';

// TensorFlow.js imports
import * as tf from '@tensorflow/tfjs';
import '@tensorflow/tfjs-backend-webgpu'; // Only WebGPU is supported
import {WebGPUBackend} from '@tensorflow/tfjs-backend-webgpu';

async function main() {
  // Initialize TensorFlow.js WebGPU backend
  await tf.setBackend('webgpu');

  // Initialize LiteRT.js's Wasm files
  await loadLiteRt('your/path/to/wasm/');

  // Make LiteRt use the same GPU device as TFJS (for tensor conversion)
  const backend = tf.backend() as WebGPUBackend;
  setWebGpuDevice(backend.device);
  // ...
}

main();

טוענים את מודל LiteRT שהומר:

const model = await loadAndCompile('path_to_model.tflite', {
  accelerator: 'webgpu', // or 'wasm'
});

כתיבת צינור עיבוד הנתונים של המודל

כותבים את הלוגיקה של העיבוד המקדים והעיבוד שאחרי העיבוד שמקשרת את המודל לאפליקציה. מומלץ להשתמש ב-TensorFlow.js לעיבוד המקדים ולעיבוד שאחרי העיבוד, אבל אם הלוגיקה לא כתובה ב-TensorFlow.js, אפשר לקרוא ל-await tensor.data כדי לקבל את הערך כ-ArrayBuffer או ל-await tensor.array כדי לקבל מערך JS מובנה.

הדוגמה הבאה היא צינור עיבוד נתונים מקצה לקצה של ResNet18:

// Wrap in a tf.tidy call to automatically clean up intermediate TensorFlow.js tensors.
// (Note: tidy only supports synchronous functions).
const top5 = tf.tidy(() => {
  // Get RGB data values from an image element and convert it to range [0, 1).
  const image = tf.browser.fromPixels(dogs, 3).div(255);

  // These preprocessing steps come from https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py#L315
  // The mean and standard deviation for the image normalization come from https://github.com/pytorch/vision/blob/main/torchvision/transforms/_presets.py#L38
  const imageData = image.resizeBilinear([224, 224])
    .sub([0.485, 0.456, 0.406])
    .div([0.229, 0.224, 0.225])
    .reshape([1, 224, 224, 3])
    .transpose([0, 3, 1, 2]);

  // You can pass inputs as a single tensor, an array, or a JS Object
  // where keys are the tensor names in the TFLite model.
  // When passing an Object, the output is also an Object.
  // Here, we're passing a single tensor, so the output is an array.
  const probabilities = runWithTfjsTensors(model, imageData)[0];

  // Get the top five classes.
  return tf.topk(probabilities, 5);
});

const values = await top5.values.data();
const indices = await top5.indices.data();
top5.values.dispose(); // Clean up the tfjs tensors.
top5.indices.dispose();

// Print the top five classes.
const classes = ... // Class names are loaded from a JSON file in the demo.
for (let i = 0; i < 5; ++i) {
  const text = `${classes[indices[i]]}: ${values[i]}`;
  console.log(text);
}

בדיקה ופתרון בעיות

בקטעים הבאים מוסבר איך לבדוק את האפליקציה ולטפל בשגיאות.

בדיקה באמצעות קלט מזויף

אחרי טעינת המודל, מומלץ לבדוק אותו באמצעות קלט פיקטיבי. כך תוכלו לזהות שגיאות בזמן ריצה לפני שתשקיעו זמן בכתיבת הלוגיקה של העיבוד המקדים והעיבוד הבא של צינור הנתונים של המודל. כדי לבדוק את זה, אפשר להשתמש בבודק המודלים LiteRT.js או לבדוק את זה באופן ידני.

LiteRT.js Model Tester

הכלי LiteRT.js Model Tester מריץ את המודל ב-GPU וב-CPU באמצעות קלטים אקראיים כדי לוודא שהמודל פועל בצורה תקינה ב-GPU. הכלי בודק את הדברים הבאים:

  • האם סוגי הנתונים של הקלט והפלט נתמכים.
  • האם כל הפעולות זמינות ב-GPU.
  • עד כמה הפלט של ה-GPU תואם לפלט של ה-CPU.
  • הביצועים של הסקת מסקנות ב-GPU.

כדי להפעיל את הכלי לבדיקת מודלים LiteRT.js, מריצים את הפקודה npm i @litertjs/model-tester ואז את הפקודה npx model-tester. תיפתח כרטיסייה בדפדפן שבה תוכלו להריץ את המודל.

בדיקה ידנית של מודלים

אם אתם מעדיפים לבדוק את המודל באופן ידני במקום להשתמש בכלי לבדיקת מודלים של LiteRT.js‏ (@litertjs/model-tester), אתם יכולים ליצור קלט מזויף ולהריץ את המודל באמצעות runWithTfjsTensors.

כדי ליצור קלט מזויף, צריך לדעת את השמות והצורות של טנסורי הקלט. אפשר למצוא אותם באמצעות LiteRT.js על ידי קריאה ל-model.getInputDetails או ל-model.getOutputDetails. דרך פשוטה למצוא אותם היא להגדיר נקודת עצירה אחרי יצירת המודל. אפשר גם להשתמש בדוח הביצועים של המודל.

אחרי שמקבלים את הצורות והשמות של הקלט והפלט, אפשר לבדוק את המודל באמצעות קלט מזויף. כך אפשר להיות בטוחים שהמודל יפעל לפני שכותבים את שאר צינור הנתונים של למידת המכונה. כך תוכלו לבדוק שכל הפעולות של המודל נתמכות. לדוגמה:

// Imports, initialization, and model loading...
// Create fake inputs for the model
const fakeInputs = model.getInputDetails().map(
    ({shape, dtype}) => tf.ones(shape, dtype));

// Run the model
const outputs = runWithTfjsTensors(model, fakeInputs);
console.log(outputs);

סוגי שגיאות

יכול להיות שחלק מהמודלים של LiteRT לא נתמכים על ידי LiteRT.js. בדרך כלל השגיאות נכללות בקטגוריות הבאות:

  • אי התאמה של צורות: באג מוכר שמשפיע רק על GPU.
  • הפעולה לא נתמכת: סביבת זמן הריצה לא תומכת בפעולה במודל. הכיסוי של קצה העורף של WebGPU מוגבל יותר מזה של המעבד (CPU), ולכן אם השגיאה הזו מופיעה ב-GPU, יכול להיות שאפשר להריץ את המודל במעבד במקום זאת.
  • סוג טנסור לא נתמך: LiteRT.js תומך רק בטנסורים מסוג int32 ו-float32 לנתוני קלט ופלט של מודלים.
  • המודל גדול מדי: יש מגבלה על הגודל של המודלים שאפשר לטעון ב-LiteRT.js.

הפעולה לא נתמכת

ההודעה הזו מציינת שהעורף שבו נעשה שימוש לא תומך באחת מהפעולות במודל. כדי להימנע מהפעולה הזו, תצטרכו לשכתב את מודל PyTorch המקורי ולהמיר אותו מחדש, או שתוכלו להריץ את המודל ב-CPU.

במקרה של BROADCAST_TO, יכול להיות שאפשר לפתור את הבעיה על ידי הגדרת הממד של חבילת הנתונים כך שיהיה זהה לכל טנסור קלט של המודל. מקרים אחרים עשויים להיות מורכבים יותר.

סוג טנסור לא נתמך

‫LiteRT.js תומך רק בטנסורים מסוג int32 ו-float32 עבור הקלט והפלט של המודל.

המודל גדול מדי

בדרך כלל השגיאה הזו מופיעה כקריאה אל Aborted() או ככשל בהקצאת זיכרון בזמן טעינת המודל. הגודל של המודלים ש-LiteRT.js יכול לטעון מוגבל, ולכן אם אתם רואים את ההודעה הזו, יכול להיות שהמודל שלכם גדול מדי. אפשר לנסות לכמת את המשקלים באמצעות ai-edge-quantizer, אבל להשאיר את החישובים בפורמט float32 או float16, ואת הקלט והפלט של המודל בפורמט float32 או int32.