在 Google Tensor 上實作 LiteRT

1. 總覽

Google Tensor SDK 用於編譯 Pixel 裝置的 LiteRT 模型。編譯後的模型可部署在 Pixel 裝置上，提升機器學習推論效能。如要使用 SDK，請先將模型轉換為 LiteRT (tflite) 模型。

本程式碼研究室以 GitHub 上的通用 Colab 為基礎，即 LiteRT AOT 編譯教學課程 Colab：在 Colab 中開啟。

目標

瞭解如何使用 LiteRT AOT (預先) 編譯器，將 TFLite 模型中的自拍分割模型編譯為 LiteRT 模型，並針對裝置端 Edge TPU 進行最佳化和編譯。

這個 Colab 也會逐步說明如何使用 Play 裝置端 AI (PODAI) 準備模型。

PODAI 可更有效率地為裝置端 AI 功能提供自訂模型。這項功能可簡化啟動、指定目標、版本管理及下載 AI 模型的程序。搭配 LiteRT EdgeTPU AOT 編譯功能，開發人員就能為各種裝置提供已編譯的 ML 模型，不必瞭解使用者的手機搭載哪些 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
    或者，您可以在 Colab 筆記本中執行這項操作：
    %env GOOGLE_TENSOR_SDK_BETA=/path/to/downloaded/compiler
  4. 然後執行下列指令來安裝套件：
    pip install ai-edge-litert-sdk-google-tensor

3. 安裝必要套件

首先，請安裝必要套件，包括含有 EdgeTPU AOT 編譯器的 ai-edge-litert-nightly，以及用於模型轉換的其他程式庫。

使用這個套件安裝 Google Tensor 的 LiteRT 後端：ai-edge-litert-sdk-google-tensor。

安裝套件後，請重新啟動工作階段，然後繼續執行安裝步驟。請勿重複安裝。

如果您打算在系統上執行設定，建議使用 Python 虛擬環境 (venv)，並在虛擬環境中執行這些指令。

解除安裝特定套件

在此之前，請先解除安裝 Colab 執行階段預設隨附的 TensorFlow。

pip uninstall -y tensorflow ai-edge-litert

安裝所有程式庫

安裝 Google Tensor 的 LiteRT 後端

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) 模型。

從 TFLite 模型編譯 EdgeTPU

這個步驟需要 TFLite 模型。如果沒有 TFLite 模型，請將模型轉換為 TFLite 格式。

取得 TFLite 模型

我們使用 MediaPipe MultiClass Segmentation 模型處理這個使用案例。

您可以在 MediaPipe 影像分割頁面取得 TFLite 模型。

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)

使用 LiteRT Python API 快速驗證 TfLite 模型

在下列範例中，您會看到遮罩圖片和混合結果。

# 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 編譯。

我們使用 ai_edge_litert.aot 的 API 編譯模型。

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 方法將所有模型匯出至磁碟。

根據預設，模型會以扁平結構儲存在輸出目錄中，每個模型名稱都會加上後端 ID 後綴。

例如：

模型檔案名稱

後端

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 上匯出及驗證

編譯完成後，請在 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，並透過 Google Play On-Device AI (PODAI) 架構傳送至使用者裝置。

AiEdgeLiteRT AOT (預先) 模組提供專為此用途設計的 ai_pack 公用程式。這些公用程式會建立 AI Pack，這是重要的資料資產。AI 套件會將編譯的模型與指定裝置設定綁在一起，確保正確的模型和素材資源能傳送到適當的使用者裝置。這對 NPU (神經網路處理單元) 編譯作業尤其重要，因為這可確保針對特定系統單晶片 (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. 設定進階選項

特定裝置或 EdgeTPU 的 NPU 編譯

根據預設，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 影像分割範例」。