Google uses AI technology to translate content into your preferred language. AI translations can contain errors.

Google Tensor での LiteRT の実装

1. 概要

Google Tensor SDK は、Pixel デバイス用の LiteRT モデルをコンパイルするために使用されます。コンパイルされたモデルは、ML 推論のパフォーマンスを向上させるために Pixel デバイスにデプロイできます。SDK を使用するには、まずモデルを LiteRT（tflite）モデルに変換する必要があります。

この Codelab は、GitHub の一般的な Colab（LiteRT AOT Compilation Tutorial Colab:）に基づいています。

目標

LiteRT AOT （Ahead-of-Time）コンパイラを使用して、TFLite モデルから、オンデバイス EdgeTPU 用に最適化およびコンパイルされた LiteRT モデルにセルフィーセグメンテーションモデルをコンパイルする方法を学びます。

この Colab では、Play for On-device AI（PODAI）でモデルを準備する手順についても説明します。

PODAI は、オンデバイス AI 機能用のカスタムモデルをより効率的に提供します。AI モデルの起動、ターゲティング、バージョン管理、ダウンロードのプロセスを簡素化します。LiteRT EdgeTPU AOT コンパイルと組み合わせることで、デベロッパーは、エンドユーザーのスマートフォンに搭載されている EdgeTPU を必知事項とすることなく、さまざまなデバイス向けにコンパイル済みの ML モデルを提供できます。

使用したモデル

使用しているモデルは、もともと MediaPipe 画像セグメンテーションガイドで公開されています。この Codelab で使用するモデルの詳細を以下に示します。

SelfieMulticlass: 人物の画像を撮影し、髪、肌、衣服などの領域を特定して、これらのアイテムの画像セグメンテーションマップを出力する LiteRT モデル。

2. 始める

Google Tensor SDK にアクセスして使用を開始する手順は次のとおりです。

登録して Google Tensor SDK にアクセスします。先に進む前に、コンパイラプラグインのダウンロードリンクが記載された Google からのメールが届くまでお待ちください。
コンパイラプラグイン（litert_plugin_compiler.tar.gz）をダウンロードして、任意のフォルダに配置します。
環境変数 GOOGLE_TENSOR_SDK_BETA を、ダウンロードしたファイルのローカルシステムパスに設定します。
bash ターミナルで次のコマンドを実行できます。
```
export GOOGLE_TENSOR_SDK_BETA=/path/to/downloaded/compiler
```
または、Colab ノートブックで次のコマンドを実行できます。
```
%env GOOGLE_TENSOR_SDK_BETA=/path/to/downloaded/compiler
```
次に、次のコマンドを実行してパッケージをインストールします。
```
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

注: ai-edge-litert-nightly ではなく安定版リリースパッケージ ai-edge-litert を使用している場合は、すべての import ステートメントの前に（スクリプトの実行ごとに）環境変数 GOOGLE_TENSOR_BACKEND_ENABLED=1 を設定する必要があります。同じ Python コードスニペットは次のとおりです。

import os
os.environ['GOOGLE_TENSOR_BACKEND_ENABLED'] = '1'

または（シェル）

export GOOGLE_TENSOR_BACKEND_ENABLED=1

5. LiteRT モデルをコンパイルする

このセクションでは、LiteRT（TFLite）モデルを直接コンパイルするなど、高度な使用方法について説明します。

TFLite モデルからの EdgeTPU コンパイル

この手順では、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)

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

EdgeTPU AOT コンパイルを使用して LiteRT モデルに変換する。

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（Ahead-of-Time）モジュールには、この目的専用の ai_pack ユーティリティが用意されています。これらのユーティリティは、重要なデータアセットである AI パック を作成します。AI パックは、コンパイル済みのモデルとデバイスターゲティング構成をバンドルし、適切なモデルとアセットが適切なユーザーデバイスに配信されるようにします。これは、特定の System-on-Chip（SoC）用に最適化されたモデルが、その SoC を搭載したデバイスにのみ到達するようにするため、NPU（Neural Processing Unit）コンパイルでは特に重要です。

# 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 画像セグメンテーションサンプルをご覧ください。