הטמעה של LiteRT ב-Google Tensor

1. סקירה כללית

ה-SDK של Google Tensor משמש להידור מודלים של LiteRT למכשירי Pixel. אפשר לפרוס את המודלים המהודרים במכשירי Pixel כדי לשפר את ביצועי ההסקה של למידת מכונה. כדי להשתמש ב-SDK, קודם צריך להמיר את המודל למודל LiteRT ‏ (tflite).

ה-Codelab הזה מבוסס על ה-Colab הכללי ב-GitHub, LiteRT AOT Compilation Tutorial Colab:פתיחה ב-Colab.

מטרה

במאמר הזה נסביר איך להשתמש בקומפיילר LiteRT AOT (קומפילציה מראש) כדי לקמפל מודל של פילוח תמונות סלפי ממודל TFLite למודל LiteRT שעבר אופטימיזציה וקומפילציה לשימוש ב-EdgeTPU במכשיר.

ב-Colab הזה מוסבר גם איך להכין מודלים באמצעות Play for On-device AI‏ (PODAI).

‫PODAI מספק מודלים מותאמים אישית לתכונות AI במכשיר בצורה יעילה יותר. הוא מפשט את התהליך של השקת מודלים של AI, טירגוט שלהם, ניהול גרסאות שלהם והורדה שלהם. בשילוב עם קומפילציה של LiteRT EdgeTPU AOT, הוא מאפשר למפתחים לספק מודלים של למידת מכונה שעברו קומפילציה למכשירים שונים, בלי שהם צריכים לדעת אילו EdgeTPU מכיל הטלפון של משתמש הקצה.

המודלים שבהם נעשה שימוש

המודלים שבהם אנחנו משתמשים פורסמו במקור במדריך לחלוקת תמונות למקטעים של MediaPipe. ריכזנו כאן פרטים על המודל שבו נעשה שימוש בסדנת הקוד הזו:

  • ‫SelfieMulticlass: מודל LiteRT שמקבל תמונה של אדם, מאתר אזורים כמו שיער, עור ובגדים, ומפיק מפת חלוקת תמונות למקטעים של הפריטים האלה.

2. שנתחיל?

כדי לקבל גישה ל-Google Tensor SDK ולהתחיל לעבוד איתו, צריך לפעול לפי השלבים הבאים:

  1. נרשמים כדי לקבל גישה ל-Google Tensor SDK. לפני שממשיכים, צריך לחכות לאימייל מ-Google עם קישור להורדה של תוסף הקומפיילר.
  2. מורידים את תוסף הקומפיילר (litert_plugin_compiler.tar.gz) וממקמים אותו בתיקייה לבחירתכם.
  3. מגדירים את משתנה הסביבה לנתיב המערכת המקומית של הקובץ שהורד, GOOGLE_TENSOR_SDK_BETA.
    אפשר להריץ את הפקודה הזו במסוף Bash:
    export GOOGLE_TENSOR_SDK_BETA=/path/to/downloaded/compiler
    אפשר גם להריץ את הפקודה הבאה בקובץ ה-notebook של Colab:
    %env GOOGLE_TENSOR_SDK_BETA=/path/to/downloaded/compiler
  4. לאחר מכן מריצים את הפקודה הבאה כדי להתקין את החבילה:
    pip install ai-edge-litert-sdk-google-tensor

3. התקנת החבילות הנדרשות

מתחילים בהתקנת החבילות הנדרשות, כולל ai-edge-litert-nightly, שמכילה את מהדר ה-AOT של EdgeTPU, וספריות אחרות שמשמשות להמרת מודלים.

משתמשים בחבילה הזו כדי להתקין את ה-backend של LiteRT ל-Google Tensor: ai-edge-litert-sdk-google-tensor.

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

אם אתם מתכננים לבצע את ההגדרה במערכת שלכם, מומלץ להשתמש בסביבה וירטואלית של Python ‏ (venv) ולהריץ את הפקודות האלה בתוך הסביבה הווירטואלית.

הסרת חבילות מסוימות

לפני כן, צריך להסיר את tensorflow שמגיע עם סביבת זמן הריצה של Colab כברירת מחדל.

pip uninstall -y tensorflow ai-edge-litert

התקנה של כל הספריות

התקנת ה-Backend של LiteRT ל-Google Tensor

pip install ai-edge-litert-sdk-google-tensor

התקנת החבילות שנותרו

pip install matplotlib huggingface-hub ai-edge-litert-nightly

4. ייבוא כל הספריות

אחרי שההתקנה מסתיימת, ממשיכים להפעלה הראשית.

מייבאים את החבילות הנדרשות:

import os
import shutil

from ai_edge_litert.aot import aot_compile as aot_lib
from ai_edge_litert.aot.ai_pack import export_lib as ai_pack_export
from ai_edge_litert.aot.vendors.google_tensor import target as gt_target
import huggingface_hub
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
import requests

5. הידור של מודל LiteRT

בקטע הזה נסביר על שימושים מתקדמים, כמו קומפילציה ישירה של מודל LiteRT ‏ (TFLite).

הידור EdgeTPU ממודל TFLite

בשלב הזה נדרש מודל TFLite. אם אין לכם מודל TFLite, אתם צריכים להמיר את המודל לפורמט TFLite.

קבלת מודל TFLite

בתרחיש לדוגמה הזה אנחנו משתמשים במודל MediaPipe MultiClass Segmentation.

מודל TFLite זמין בדף חלוקת תמונות למקטעים ב-MediaPipe.

work_dir = '.'

model_url = 'https://storage.googleapis.com/mediapipe-models/image_segmenter/selfie_multiclass_256x256/float32/latest/selfie_multiclass_256x256.tflite'
tflite_model_path = os.path.join(work_dir, 'selfie_multiclass_256x256.tflite')

model_content = requests.get(model_url)

with open(tflite_model_path, 'wb') as fout:
  fout.write(model_content.content)

אימות מהיר של מודל TfLite באמצעות API בשפת Python של LiteRT

בדוגמה הבאה אפשר לראות גם את תמונת המסכה וגם את התוצאה הממוזגת.

# Downloading Testing image

test_image = huggingface_hub.hf_hub_download(
    repo_id="litert-community/MediaPipe-Selfie-Segmentation",
    filename="test_img.png",
)
pil_image = Image.open(test_image).convert("RGB").resize((256, 256))

from ai_edge_litert.compiled_model import CompiledModel

SEGMENT_COLORS = [
    (0, 0, 0),
    (255, 0, 0),
    (0, 255, 0),
    (0, 0, 255),
    (255, 255, 0),
    (255, 0, 255),
]
INPUT_SIZE = (256, 256)
NUM_CLASSES = 6

# Load the model and image
model = CompiledModel.from_file(tflite_model_path)
original_image = np.array(Image.open(test_image).convert('RGB'))
img_array = np.array(pil_image).astype(np.float32)

# Normalize the image
normalized = (img_array - 127.5) / 127.5
normalized = np.ascontiguousarray(normalized, dtype=np.float32)

# Run inference
sig_idx = 0
input_buffers = model.create_input_buffers(sig_idx)
output_buffers = model.create_output_buffers(sig_idx)
input_data = normalized.reshape(-1)
input_buffers[0].write(input_data)
model.run_by_index(sig_idx, input_buffers, output_buffers)

# Get output data
height, width = INPUT_SIZE
output_size = height * width * NUM_CLASSES
output_data = output_buffers[0].read(output_size, np.float32)
output_data = output_data.reshape(height, width, NUM_CLASSES)
mask = np.argmax(output_data, axis=2).astype(np.uint8)

# Create colored mask
colored_mask = np.zeros((height, width, 3), dtype=np.uint8)
for label_idx in range(NUM_CLASSES):
  class_mask = mask == label_idx
  color = SEGMENT_COLORS[label_idx]
  colored_mask[class_mask] = color

# Blend with original image
# Resize colored mask to match original image if necessary
if original_image.shape[:2] != colored_mask.shape[:2]:
  colored_mask_pil = Image.fromarray(colored_mask)
  colored_mask_pil = colored_mask_pil.resize(
      (original_image.shape[1], original_image.shape[0])
  )
  colored_mask = np.array(colored_mask_pil)

# Blend images with alpha 0.5
alpha = 0.5
blended_image = (
    original_image * (1 - alpha) + colored_mask * alpha
).astype(np.uint8)

# Display them
fig, axes = plt.subplots(1, 3, figsize=(9, 3))

for idx, (title, image) in enumerate([
    ('Original Image', original_image),
    ('Colored Mask', colored_mask),
    ('Blended Image', blended_image),
]):
  axes[idx].imshow(image)
  axes[idx].set_title(title)
  axes[idx].axis('off')

plt.tight_layout()
plt.show()

המרת המודל ל-LiteRT, עם קומפילציה של EdgeTPU AOT.

אנחנו משתמשים בממשקי ה-API של ai_edge_litert.aot כדי להרכיב את המודל.

compiled_models = aot_lib.aot_compile(tflite_model_path, keep_going=True)

# This variable will be used later to create the AI Pack.
all_google_tensor_compiled_models = compiled_models

# Print Compilation Report
print(all_google_tensor_compiled_models.compilation_report())

# Saving compiled models to disk. This saves all the compiled models, and a CPU
# fallback model.
all_google_tensor_compiled_models.export(
    work_dir, model_name='selfie_segmentation'
)

בסיום ההידור, משתמשים בשיטה model.export כדי לייצא את כל המודלים לדיסק.

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

לדוגמה:

שם הקובץ של הדגם

בק-אנד

SoC

הערה

selfie_segmentation_fallback.tflite

CPU/GPU

לא רלוונטי

לא רלוונטי

selfie_segmentation_Google_Tensor_G3.tflite

Google

Tensor_G3

Google Tensor G3

selfie_segmentation_Google_Tensor_G4.tflite

Google

Tensor_G4

Google Tensor G4

selfie_segmentation_Google_Tensor_G5.tflite

Google

Tensor_G5

Google Tensor G5

6. ייצוא ואימות ב-CPU

אחרי שהקומפילציה מסתיימת, מאמתים את מודל TFLite במעבד. כדי לעשות זאת, משתמשים ב'מודל חלופי' שנוצר במהלך ההידור.

# Run LiteRT with test image
from ai_edge_litert.compiled_model import CompiledModel

# Normalize the image to [-1, 1]
img_array = np.array(pil_image, dtype=np.float32)
normalized = (img_array - 127.5) / 127.5
numpy_array = np.ascontiguousarray(normalized)[None, ...]

cpu_model_path = os.path.join(work_dir, "selfie_segmentation_fallback.tflite")
cm_model = CompiledModel.from_file(cpu_model_path)
sig_idx = 0
input_buffers = cm_model.create_input_buffers(sig_idx)
output_buffers = cm_model.create_output_buffers(sig_idx)
input_buffers[0].write(numpy_array)
cm_model.run_by_index(sig_idx, input_buffers, output_buffers)

# Read the 6-channel output and apply argmax
output_data = output_buffers[0].read(256 * 256 * 6, np.float32)
output_data = output_data.reshape((256, 256, 6))
mask = np.argmax(output_data, axis=2).astype(np.uint8)

# Create a colored mask using the previously defined SEGMENT_COLORS
colored_mask = np.zeros((256, 256, 3), dtype=np.uint8)
for label_idx in range(6):
  class_mask = mask == label_idx
  color = SEGMENT_COLORS[label_idx]
  colored_mask[class_mask] = color

mask_image = Image.fromarray(colored_mask)

# Show output results
fig, axes = plt.subplots(1, 2, figsize=(9, 3))

for idx, (title, image) in enumerate([
    ('Test Image', pil_image),
    ('TFLite Mask Image', mask_image),
]):
  axes[idx].imshow(image)
  axes[idx].set_title(title)
  axes[idx].axis('off')

plt.tight_layout()
plt.show()

7. ייצוא מודלים ל-PODAI

אחרי שמאמתים את המודלים, השלב החשוב הבא הוא להכין אותם לפריסה. בקטע הזה מוסבר איך לארוז את המודלים שעברו קומפילציה כדי להעלות אותם ל-Google Play, וכך לאפשר את ההפצה שלהם למכשירי משתמשים באמצעות מסגרת ה-AI במכשיר (PODAI) של Google Play.

מודול ה-AOT ‏ (Ahead-of-Time) של AiEdgeLiteRT מספק ai_pack כלי עזר במיוחד למטרה הזו. הכלים האלה יוצרים חבילת AI, שהיא נכס נתונים חיוני. חבילת AI כוללת את המודלים המהודרים שלכם עם הגדרות טירגוט למכשירים, כדי לוודא שהמודלים והנכסים הנכונים מועברים למכשירי המשתמשים המתאימים. זה חשוב במיוחד לקומפילציות של NPU (יחידת עיבוד עצבי), כי זה מבטיח שמודלים שעברו אופטימיזציה למערכת ספציפית על שבב (SoC) יגיעו רק למכשירים שמצוידים במערכת הזו על שבב.

# Configuring the AI Pack
os.makedirs('selfie_multiclass', exist_ok=True)
ai_pack_dir = os.path.join(work_dir, 'ai_pack')
ai_pack_name = 'selfie_segmentation'
litert_model_name = 'segmentation_model'

# Clean up
shutil.rmtree(ai_pack_dir, ignore_errors=True)

# Export
ai_pack_export.export(
    all_google_tensor_compiled_models,
    ai_pack_dir,
    ai_pack_name,
    litert_model_name
)

בדיקת המקור של חבילת AI

def list_files(startpath):
  """Function to print out the tree structure of a directory."""
  for root, dirs, files in os.walk(startpath):
    level = root.replace(startpath, '').count(os.sep)
    indent = ' ' * 4 * (level)
    print('{}{}/'.format(indent, os.path.basename(root)))
    subindent = ' ' * 4 * (level + 1)
    for f in files:
      print('{}{}'.format(subindent, f))
"""View the files generated within the AI pack directory"""
list_files(ai_pack_dir)

8. הגדרה של אפשרויות מתקדמות

קומפילציה של NPU למכשיר ספציפי או ל-EdgeTPU

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

הדוגמה הבאה עוברת קומפילציה ל-Google Tensor G5.

# Specifying the compilation target
tensor_g5_target = gt_target.Target(gt_target.SocModel.TENSOR_G5)

# Compile from the TFLite model for a specific target
compiled_models = aot_lib.aot_compile(
    tflite_model_path,
    target=[tensor_g5_target],
    keep_going=False,  # We want to error out when there's failure.
)

print(compiled_models.compilation_report())

דגלים של הידור ל-Google Tensor

התאמה אישית של תהליך ההידור באמצעות דגלי הידור. בדוגמה הזו נעשה שימוש בדגל הזה: google_tensor_truncation_type="half"

כשמקמפלים מודל TFLite

compiled_models = aot_lib.aot_compile(
    tflite_model_path,
    target=[tensor_g5_target],
    keep_going=False,
    google_tensor_truncation_type="half"
)

9. השלבים הבאים

מזל טוב!

המודלים מוכנים לשימוש ב-PODAI.

עכשיו עוברים אל Android Studio כדי לבצע את השלבים הבאים. פרטים נוספים זמינים במאמר דוגמאות לחלוקת תמונות למקטעים ב-LiteRT.