การใช้งาน LiteRT ใน Google Tensor

1. ภาพรวม

Google Tensor SDK ใช้เพื่อคอมไพล์โมเดล LiteRT สำหรับอุปกรณ์ Pixel และสามารถนำโมเดลที่คอมไพล์แล้วไปใช้งานในอุปกรณ์ Pixel เพื่อเพิ่มประสิทธิภาพการอนุมาน ML ได้ หากต้องการใช้ SDK คุณต้องแปลงโมเดลเป็นโมเดล LiteRT (tflite) ก่อน

Codelab นี้ อิงตาม Colab ทั่วไปใน GitHub LiteRT AOT Compilation Tutorial Colab:เปิดใน Colab

วัตถุประสงค์

เรียนรู้วิธีใช้คอมไพเลอร์ LiteRT AOT (Ahead-of-Time) เพื่อคอมไพล์ โมเดลการแบ่งกลุ่มเซลฟี จากโมเดล TFLite เป็นโมเดล LiteRT ที่มีการเพิ่มประสิทธิภาพและคอมไพล์สำหรับ EdgeTPU ในอุปกรณ์

นอกจากนี้ Colab นี้ยังแนะนำขั้นตอนการเตรียมโมเดลด้วย Play for On-device AI (PODAI)

PODAI ช่วยให้ส่งโมเดลที่กำหนดเองสำหรับฟีเจอร์ AI ในอุปกรณ์ได้อย่างมีประสิทธิภาพมากขึ้น และช่วยลดความซับซ้อนของกระบวนการเปิดตัว กำหนดเป้าหมาย การกำหนดเวอร์ชัน และดาวน์โหลดโมเดล AI เมื่อใช้ร่วมกับการคอมไพล์ LiteRT EdgeTPU AOT จะช่วยให้นักพัฒนาแอปส่งโมเดล ML ที่คอมไพล์แล้วสำหรับอุปกรณ์ต่างๆ ได้โดยไม่ต้องทราบว่าโทรศัพท์ของผู้ใช้ปลายทางมี EdgeTPU รุ่นใด

โมเดลที่ใช้

โมเดลที่เราใช้เผยแพร่ครั้งแรกในคู่มือการแบ่งกลุ่มรูปภาพของ MediaPipe รายละเอียดบางส่วนเกี่ยวกับโมเดลที่ใช้ใน Codelab นี้มีดังนี้

  • 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
    
    หรือเรียกใช้คำสั่งนี้ในสมุดบันทึก Colab
    %env GOOGLE_TENSOR_SDK_BETA=/path/to/downloaded/compiler
    
  4. จากนั้นเรียกใช้คำสั่งนี้เพื่อติดตั้งแพ็กเกจ
    pip install ai-edge-litert-sdk-google-tensor
    

3. ติดตั้งแพ็กเกจที่จำเป็น

เริ่มต้นด้วยการติดตั้งแพ็กเกจที่จำเป็น ซึ่งรวมถึง ai-edge-litert-nightly ที่มีคอมไพเลอร์ EdgeTPU AOT และไลบรารีอื่นๆ ที่คุณใช้สำหรับการแปลงโมเดล

ใช้แพ็กเกจนี้เพื่อติดตั้งแบ็กเอนด์ LiteRT สำหรับ Google Tensor: ai-edge-litert-sdk-google-tensor

หลังจากติดตั้งแพ็กเกจแล้ว ให้รีสตาร์ทเซสชันและดำเนินการต่อจากขั้นตอนการติดตั้ง อย่า ติดตั้งซ้ำ

หากคุณวางแผนที่จะทำการตั้งค่าในระบบ เราขอแนะนำให้คุณใช้สภาพแวดล้อมเสมือน (venv) ของ Python และเรียกใช้คำสั่งเหล่านี้ภายในสภาพแวดล้อมเสมือน

ถอนการติดตั้งแพ็กเกจบางรายการ

ก่อนอื่น ให้ถอนการติดตั้ง TensorFlow ที่มาพร้อมกับรันไทม์ Colab โดยค่าเริ่มต้น

pip uninstall -y tensorflow ai-edge-litert

ติดตั้งไลบรารีทั้งหมด

ติดตั้งแบ็กเอนด์ 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 สำหรับกรณีการใช้งานนี้

โมเดล 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 อย่างรวดเร็วโดยใช้ LiteRT Python API

ในตัวอย่างต่อไปนี้ คุณจะเห็นทั้งรูปภาพมาสก์และผลลัพธ์ที่ผสานรวม

# 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 ใน CPU โดยใช้ "โมเดลสำรอง" ที่สร้างขึ้นระหว่างการคอมไพล์

# 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 ซึ่งจะช่วยให้ส่งโมเดลไปยังอุปกรณ์ของผู้ใช้ผ่านเฟรมเวิร์ก Google Play On-Device AI (PODAI) ได้

โมดูล AiEdgeLiteRT AOT (Ahead-of-Time) มีเครื่องมือ ai_pack สำหรับวัตถุประสงค์นี้โดยเฉพาะ เครื่องมือเหล่านี้จะสร้างแพ็ก AI ซึ่งเป็นชิ้นงานข้อมูลที่สำคัญ แพ็ก AI จะรวมโมเดลที่คอมไพล์แล้วเข้ากับการกำหนดค่าการกำหนดเป้าหมายอุปกรณ์ เพื่อให้มั่นใจว่าระบบจะส่งโมเดลและชิ้นงานที่ถูกต้องไปยังอุปกรณ์ของผู้ใช้ที่เหมาะสม ซึ่งมีความสำคัญอย่างยิ่งสำหรับการคอมไพล์ NPU (หน่วยประมวลผลระบบประสาท) เนื่องจากจะช่วยให้มั่นใจว่าโมเดลที่เพิ่มประสิทธิภาพสำหรับ System-on-Chip (SoC) ที่เฉพาะเจาะจงจะไปถึงเฉพาะอุปกรณ์ที่มี 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