Menerapkan LiteRT di Google Tensor

1. Ringkasan

Google Tensor SDK digunakan untuk mengompilasi model LiteRT untuk perangkat Pixel. Model yang dikompilasi dapat di-deploy di perangkat Pixel untuk meningkatkan performa inferensi ML. Untuk menggunakan SDK, Anda harus mengonversi model menjadi model LiteRT (tflite).

Codelab ini berdasarkan colab umum di GitHub, LiteRT AOT Compilation Tutorial Colab:Buka di Colab.

Tujuan

Pelajari cara menggunakan compiler LiteRT AOT (ahead of time) untuk mengompilasi model segmentasi selfie dari model TFLite menjadi model LiteRT yang dioptimalkan dan dikompilasi untuk EdgeTPU di perangkat.

Colab ini juga memandu Anda melalui langkah-langkah untuk menyiapkan model dengan Play for On-device AI (PODAI).

PODAI menghadirkan model kustom untuk fitur AI di perangkat secara lebih efisien. Fitur ini menyederhanakan proses peluncuran, penargetan, pembuatan versi, dan download model AI Anda. Jika dikombinasikan dengan kompilasi AOT EdgeTPU LiteRT, fitur ini memungkinkan developer menghadirkan model ML yang dikompilasi untuk berbagai perangkat tanpa perlu mengetahui EdgeTPU yang ada di ponsel pengguna akhir.

Model yang digunakan

Model yang kami gunakan awalnya dipublikasikan di panduan Segmentasi gambar MediaPipe. Berikut beberapa detail terkait model yang digunakan dalam codelab ini:

  • SelfieMulticlass: Model LiteRT yang mengambil gambar seseorang, menemukan area seperti rambut, kulit, dan pakaian, serta menampilkan peta segmentasi gambar untuk item ini.

2. Mulai

Ikuti langkah-langkah berikut untuk mendapatkan akses dan memulai Google Tensor SDK:

  1. Daftar untuk mendapatkan akses ke Google Tensor SDK. Sebelum melanjutkan, Anda harus menunggu email dari Google yang berisi link download untuk plugin compiler.
  2. Download plugin compiler (litert_plugin_compiler.tar.gz) dan tempatkan di folder pilihan Anda.
  3. Tetapkan variabel lingkungan ke jalur sistem lokal file yang didownload, GOOGLE_TENSOR_SDK_BETA.
    Anda dapat menjalankan perintah ini di terminal bash:
    export GOOGLE_TENSOR_SDK_BETA=/path/to/downloaded/compiler
    
    Atau Anda dapat menjalankannya di notebook Colab:
    %env GOOGLE_TENSOR_SDK_BETA=/path/to/downloaded/compiler
    
  4. Kemudian, jalankan perintah ini untuk menginstal paket:
    pip install ai-edge-litert-sdk-google-tensor
    

3. Menginstal paket yang diperlukan

Mulailah dengan menginstal paket yang diperlukan, termasuk ai-edge-litert-nightly, yang berisi compiler AOT EdgeTPU, dan library lain yang Anda gunakan untuk konversi model.

Gunakan paket ini untuk menginstal backend LiteRT untuk Google Tensor: ai-edge-litert-sdk-google-tensor.

Setelah menginstal paket, mulai ulang sesi dan lanjutkan dari langkah penginstalan. Jangan ulangi penginstalan.

Jika Anda berencana melakukan penyiapan di sistem, sebaiknya gunakan lingkungan virtual Python (venv) dan jalankan perintah ini dalam lingkungan virtual.

Meng-uninstal paket tertentu

Sebelum itu, uninstal tensorflow yang disertakan dengan runtime colab secara default.

pip uninstall -y tensorflow ai-edge-litert

Menginstal semua library

Menginstal backend LiteRT untuk Google Tensor

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

Menginstal paket yang tersisa

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

4. Mengimpor semua library

Lanjutkan ke eksekusi utama setelah penginstalan selesai.

Impor paket yang diperlukan:

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. Mengompilasi model LiteRT

Bagian ini membahas penggunaan lanjutan seperti mengompilasi model LiteRT (TFLite) secara langsung.

Kompilasi EdgeTPU dari model TFLite

Langkah ini memerlukan model TFLite. Jika tidak memiliki model TFLite, konversikan model Anda ke format TFLite.

Mendapatkan Model TFLite

Kami menggunakan model Segmentasi MultiClass MediaPipe untuk kasus penggunaan ini.

Model TFLite tersedia dari halaman Segmentasi gambar 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)

Memverifikasi model TfLite dengan cepat menggunakan LiteRT Python API

Dalam contoh berikut, Anda akan melihat gambar mask dan hasil campuran.

# 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()

Mengonversi ke model LiteRT, dengan kompilasi AOT EdgeTPU.

Kami menggunakan API dari ai_edge_litert.aot untuk mengompilasi model.

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

Setelah kompilasi selesai, gunakan metode model.export untuk mengekspor semua model ke disk.

Secara default, model disimpan dalam struktur datar di direktori output, dengan setiap nama model diberi akhiran ID backend.

Contoh:

Nama file model

Backend

SoC

Catatan

selfie_segmentation_fallback.tflite

CPU/GPU

T/A

T/A

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. Mengekspor dan memvalidasi di CPU

Setelah kompilasi selesai, verifikasi model TFLite di CPU. Lakukan hal ini menggunakan "model pengganti" yang dihasilkan selama kompilasi.

# 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. Mengekspor Model untuk PODAI

Setelah model Anda diverifikasi, langkah penting berikutnya adalah menyiapkan model untuk deployment. Bagian ini menjelaskan cara mengemas model yang dikompilasi untuk diupload ke Google Play, sehingga dapat dikirim ke perangkat pengguna melalui framework Google Play On-Device AI (PODAI).

Modul AiEdgeLiteRT AOT (Ahead-of-Time) menyediakan utilitas ai_pack khusus untuk tujuan ini. Utilitas ini membuat Paket AI, yang merupakan aset data penting. Paket AI menggabungkan model yang dikompilasi dengan konfigurasi penargetan perangkat, sehingga memastikan model dan aset yang benar dikirimkan ke perangkat pengguna yang sesuai. Hal ini sangat penting untuk kompilasi NPU (Neural Processing Unit), karena memastikan bahwa model yang dioptimalkan untuk System-on-Chip (SoC) tertentu hanya menjangkau perangkat yang dilengkapi dengan SoC tersebut.

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

Memeriksa sumber Paket 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. Mengonfigurasi opsi lanjutan

Kompilasi NPU untuk perangkat atau EdgeTPU tertentu

Secara default, kompilasi AOT LiteRT dikompilasi ke semua backend terdaftar. Untuk pengembangan lokal, Anda mungkin hanya ingin mengompilasi untuk perangkat tertentu, seperti ponsel pengembangan. Capai hal ini dengan memberikan target kompilasi secara eksplisit.

Contoh berikut mengompilasi ke 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())

Flag kompilasi untuk Google Tensor

Sesuaikan proses kompilasi melalui flag kompilasi. Di sini, flag ini digunakan: google_tensor_truncation_type="half"

Saat mengompilasi model TFLite

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

9. Langkah berikutnya

Selamat!

Model Anda siap digunakan oleh PODAI.

Sekarang, lanjutkan ke Android Studio untuk langkah-langkah berikut; lihat contoh segmentasi gambar LiteRT untuk mengetahui detailnya.