LiteRT auf Google Tensor implementieren

1. Übersicht

Das Google Tensor SDK wird verwendet, um LiteRT-Modelle für Pixel-Geräte zu kompilieren. Die kompilierten Modelle können auf Pixel-Geräten bereitgestellt werden, um die ML-Inferenzleistung zu verbessern. Damit Sie das SDK verwenden können, müssen Sie Ihr Modell zuerst in ein LiteRT-Modell (tflite) konvertieren.

Dieses Codelab basiert auf dem allgemeinen Colab auf GitHub, LiteRT AOT Compilation Tutorial Colab.

Ziel

Hier erfahren Sie, wie Sie mit dem LiteRT-AOT-Compiler (Ahead-of-Time) ein Selfie-Segmentierungsmodell aus einem TFLite-Modell in ein LiteRT-Modell kompilieren, das für EdgeTPUs auf dem Gerät optimiert und kompiliert wurde.

In diesem Colab werden auch die Schritte zum Vorbereiten von Modellen mit Play for On-device AI (PODAI) beschrieben.

Mit PODAI können benutzerdefinierte Modelle für KI-Funktionen auf dem Gerät effizienter bereitgestellt werden. Es vereinfacht das Starten, Ausrichten, Versionieren und Herunterladen Ihrer KI-Modelle. In Kombination mit der AOT-Kompilierung von LiteRT EdgeTPU können Entwickler kompilierte ML-Modelle für verschiedene Geräte bereitstellen, ohne wissen zu müssen, welche EdgeTPUs das Smartphone des Endnutzers enthält.

Verwendete Modelle

Die von uns verwendeten Modelle wurden ursprünglich im MediaPipe-Leitfaden zur Bildsegmentierung veröffentlicht. Hier sind einige Details zum im Codelab verwendeten Modell:

• SelfieMulticlass: Ein LiteRT-Modell, das ein Bild einer Person als Eingabe verwendet, Bereiche wie Haare, Haut und Kleidung lokalisiert und eine Segmentierungskarte für diese Elemente ausgibt.

2. Jetzt starten

So erhalten Sie Zugriff auf das Google Tensor SDK und können es verwenden:

1. Registrieren Sie sich, um Zugriff auf das Google Tensor SDK zu erhalten. Bevor Sie fortfahren, müssen Sie auf eine E‑Mail von Google mit dem Downloadlink für das Compiler-Plug-in warten.
2. Laden Sie das Compiler-Plug-in (litert_plugin_compiler.tar.gz) herunter und speichern Sie es in einem Ordner Ihrer Wahl.
3. Legen Sie die Umgebungsvariable auf den lokalen Systempfad der heruntergeladenen Datei fest, GOOGLE_TENSOR_SDK_BETA.
    Sie können diesen Befehl in Ihrem Bash-Terminal ausführen:

   export GOOGLE_TENSOR_SDK_BETA=/path/to/downloaded/compiler
   

   Alternativ können Sie diesen Code in Ihrem Colab-Notebook ausführen:

   %env GOOGLE_TENSOR_SDK_BETA=/path/to/downloaded/compiler
   

4. Führen Sie dann diesen Befehl aus, um das Paket zu installieren:

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

3. Erforderliche Pakete installieren

Installieren Sie zuerst die erforderlichen Pakete, einschließlich ai-edge-litert-nightly, das den AOT-Compiler für EdgeTPU enthält, und andere Bibliotheken, die Sie für die Modellkonvertierung verwenden.

Mit diesem Paket können Sie das LiteRT-Backend für Google Tensor installieren: ai-edge-litert-sdk-google-tensor.

Starten Sie die Sitzung nach der Installation der Pakete neu und fahren Sie mit den Installationsschritten fort. Wiederholen Sie die Installation nicht.

Wichtig:Wenn Sie ein Colab-Notebook oder ein Jupyter-Notebook verwenden möchten, fügen Sie vor jedem „pip install“- und „pip uninstall“-Befehl ein Ausrufezeichen ! ein. Dies ist erforderlich, wenn Sie den pip-Befehl in eine Colab-Zelle kopieren und einfügen.

Wenn Sie die Einrichtung auf Ihrem System vornehmen möchten, empfehlen wir, eine virtuelle Python-Umgebung (venv) zu verwenden und diese Befehle in der virtuellen Umgebung auszuführen.

Bestimmte Pakete deinstallieren

Deinstallieren Sie zuerst die TensorFlow-Version, die standardmäßig in der Colab-Laufzeit enthalten ist.

Hinweis:In diesem Schritt werden möglicherweise Warnungen angezeigt, dass Pakete nicht installiert sind. Sie können diese Meldungen ignorieren.

pip uninstall -y tensorflow ai-edge-litert

Alle Bibliotheken installieren

LiteRT-Backend für Google Tensor installieren

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

Verbleibende Pakete installieren

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

Wichtig:Wenn Sie ein Colab- oder Jupyter-Notebook verwenden:

• Starten Sie die Colab-Laufzeit nach Abschluss der Installation neu.
• Starten Sie die Laufzeit von Colab neu, indem Sie in der Menüleiste auf Laufzeit > Sitzung neu starten klicken, damit die installierten Pakete richtig angezeigt werden. Der Neustart sollte erst erfolgen, wenn alle Installationen abgeschlossen sind.
• Nach dem Neustart der Laufzeitumgebung führen Sie nur die folgenden Befehle aus . Sie müssen die vorherigen nicht noch einmal ausführen.

4. Alle Bibliotheken importieren

Fahren Sie mit der Hauptausführung fort, nachdem die Installation abgeschlossen ist.

Empfehlung:Verwenden Sie für diesen und die folgenden Abschnitte ein Jupyter-Notebook oder fügen Sie die Befehle einem Python-Script hinzu.

Importieren Sie die erforderlichen Pakete:

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

Hinweis:Wenn Sie das stabile Release-Paket ai-edge-litert anstelle von ai-edge-litert-nightly verwenden, müssen Sie diese Umgebungsvariable vor allen Importanweisungen festlegen (für jeden Lauf des Skripts): GOOGLE_TENSOR_BACKEND_ENABLED=1. Python-Code-Snippet für dasselbe:

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

ODER (Shell)

export GOOGLE_TENSOR_BACKEND_ENABLED=1

5. LiteRT-Modell kompilieren

In diesem Abschnitt werden erweiterte Anwendungsfälle wie das direkte Kompilieren eines LiteRT-Modells (TFLite) behandelt.

EdgeTPU-Kompilierung aus TFLite-Modell

Für diesen Schritt ist ein TFLite-Modell erforderlich. Wenn Sie kein TFLite-Modell haben, konvertieren Sie Ihr Modell in das TFLite-Format.

TFLite-Modell abrufen

Für diesen Anwendungsfall verwenden wir das Modell MediaPipe MultiClass Segmentation.

Das TFLite-Modell ist auf der Seite MediaPipe Image segmentation (MediaPipe-Bildsegmentierung) verfügbar.

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-Modell schnell mit der LiteRT Python API überprüfen

Im folgenden Beispiel sehen Sie sowohl das Maskenbild als auch das zusammengefügte Ergebnis.

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

Konvertieren Sie das Modell in ein LiteRT-Modell mit EdgeTPU-AOT-Kompilierung.

Wir verwenden die APIs aus ai_edge_litert.aot, um das Modell zu kompilieren.

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

Wenn die Kompilierung abgeschlossen ist, verwenden Sie die Methode model.export, um alle Modelle auf die Festplatte zu exportieren.

Standardmäßig werden die Modelle in einer flachen Struktur im Ausgabeverzeichnis gespeichert. An jeden Modellnamen wird die Backend-ID angehängt.

Beispiel:

Dateiname des Modells	Backend	SoC	Hinweis
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. Auf der CPU exportieren und validieren

Prüfen Sie das TFLite-Modell auf der CPU, sobald die Kompilierung abgeschlossen ist. Verwenden Sie dazu das während der Kompilierung generierte „Fallback-Modell“.

# 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. Modelle für PODAI exportieren

Nachdem Ihre Modelle überprüft wurden, müssen Sie sie für die Bereitstellung vorbereiten. In diesem Abschnitt wird beschrieben, wie Sie Ihre kompilierten Modelle für den Upload zu Google Play verpacken, damit sie über das Google Play On-Device AI-Framework (PODAI) auf Nutzergeräte übertragen werden können.

Das AOT-Modul (Ahead-of-Time) von AiEdgeLiteRT bietet ai_pack-Dienstprogramme speziell für diesen Zweck. Mit diesen Dienstprogrammen wird ein AI Pack erstellt, das ein wichtiges Daten-Asset ist. In einem AI Pack werden Ihre kompilierten Modelle mit Geräte-Targeting-Konfigurationen gebündelt. So wird dafür gesorgt, dass die richtigen Modelle und Assets auf den entsprechenden Nutzergeräten bereitgestellt werden. Dies ist besonders wichtig für NPU-Kompilierungen (Neural Processing Unit), da so sichergestellt wird, dass Modelle, die für ein bestimmtes System-on-Chip (SoC) optimiert sind, nur auf Geräten mit diesem SoC ausgeführt werden.

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

KI‑Paketquelle prüfen

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. Erweiterte Konfigurationen

NPU-Kompilierung für ein bestimmtes Gerät oder eine bestimmte EdgeTPU

Standardmäßig wird die LiteRT-AOT-Kompilierung für alle registrierten Backends kompiliert. Für die lokale Entwicklung möchten Sie möglicherweise nur für bestimmte Geräte kompilieren, z. B. für Entwickler-Smartphones. Geben Sie dazu die Kompilierungsziele explizit an.

Das folgende Beispiel wird für Google Tensor G5 kompiliert.

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

Kompilierungs-Flags für Google Tensor

Kompilierungsprozess über Kompilierungsflags anpassen Dieses Flag wird hier verwendet: google_tensor_truncation_type="half"

Beim Kompilieren eines TFLite-Modells

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

9. Nächste Schritte

Herzlichen Glückwunsch!

Ihre Modelle können von PODAI verwendet werden.

Fahren Sie nun mit Android Studio fort. Weitere Informationen finden Sie unter LiteRT-Beispiele für die Bildsegmentierung.