Wdrażanie LiteRT na Google Tensor

1. Przegląd

Pakiet Google Tensor SDK służy do kompilowania modeli LiteRT na urządzenia Pixel. Skompilowane modele można wdrażać na urządzeniach Pixel, aby zwiększyć wydajność wnioskowania ML. Aby korzystać z pakietu SDK, musisz najpierw przekonwertować model na model LiteRT (tflite).

To ćwiczenie jest oparte na ogólnym ćwiczeniu Colab na GitHubie LiteRT AOT Compilation Tutorial Colab:.

Cel

Dowiedz się, jak używać kompilatora LiteRT AOT (ahead of time) do kompilowania modelu segmentacji selfie z modelu TFLite do modelu LiteRT, który jest zoptymalizowany i skompilowany pod kątem EdgeTPU na urządzeniu.

To ćwiczenie Colab przeprowadzi Cię też przez proces przygotowywania modeli za pomocą funkcji Play for On-device AI (PODAI).

PODAI wydajniej dostarcza modele niestandardowe na potrzeby funkcji AI na urządzeniu. Upraszcza proces uruchamiania, kierowania, obsługi wersji i pobierania modeli AI. W połączeniu z kompilacją LiteRT EdgeTPU AOT umożliwia programistom dostarczanie skompilowanych modeli ML na różne urządzenia bez konieczności dostępu do informacji o tym, jakie EdgeTPU zawiera telefon użytkownika.

Używane modele

Modele, których używamy, zostały pierwotnie opublikowane w przewodniku MediaPipe segmentacja obrazu. Oto kilka szczegółów dotyczących modelu używanego w tym ćwiczeniu:

• SelfieMulticlass: model LiteRT, który pobiera obraz osoby, lokalizuje obszary takie jak włosy, skóra i ubranie oraz generuje mapę segmentacji obrazu tych elementów.

2. Rozpocznij

Aby uzyskać dostęp do pakietu Google Tensor SDK i rozpocząć z niego korzystanie:

1. Zarejestruj się, aby uzyskać dostęp do pakietu Google Tensor SDK. Zanim przejdziesz dalej, musisz poczekać na e-maila od Google z linkiem do pobrania wtyczki kompilatora.
2. Pobierz wtyczkę kompilatora (litert_plugin_compiler.tar.gz) i umieść ją w wybranym folderze.
3. Ustaw zmienną środowiskową na ścieżkę lokalną pobranego pliku, GOOGLE_TENSOR_SDK_BETA.
   To polecenie możesz uruchomić w terminalu bash:

   export GOOGLE_TENSOR_SDK_BETA=/path/to/downloaded/compiler
   

   Możesz też uruchomić je w notatniku Colab:

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

4. Następnie uruchom to polecenie, aby zainstalować pakiet:

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

3. Instalowanie wymaganych pakietów

Zacznij od zainstalowania wymaganych pakietów, w tym ai-edge-litert-nightly, który zawiera kompilator EdgeTPU AOT, oraz innych bibliotek używanych do konwersji modeli.

Użyj tego pakietu, aby zainstalować backend LiteRT dla Google Tensor: ai-edge-litert-sdk-google-tensor.

Po zainstalowaniu pakietów uruchom sesję ponownie i kontynuuj od kroków instalacji. Nie powtarzaj instalacji.

Ważne: jeśli planujesz używać notatnika Colab lub Jupyter, przed każdym poleceniem pip install i pip uninstall dodaj wykrzyknik !. Zrób to, gdy skopiujesz i wkleisz polecenie pip do komórki Colab.

Jeśli planujesz skonfigurować system, zalecamy użycie środowiska wirtualnego Python (venv) i uruchomienie tych poleceń w środowisku wirtualnym.

Odinstalowywanie niektórych pakietów

Najpierw odinstaluj TensorFlow, który domyślnie jest dostarczany ze środowiskiem wykonawczym Colab.

Uwaga: w tym kroku mogą pojawić się ostrzeżenia o nieinstalowaniu pakietów. Możesz je zignorować.

pip uninstall -y tensorflow ai-edge-litert

Instalowanie wszystkich bibliotek

Instalowanie backendu LiteRT dla Google Tensor

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

Instalowanie pozostałych pakietów

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

Ważne: jeśli używasz notatnika Colab lub Jupyter:

• Po zakończeniu instalacji uruchom ponownie środowisko wykonawcze Colab.
• Aby zainstalowane pakiety były prawidłowo odzwierciedlane, uruchom ponownie środowisko wykonawcze Colab, klikając na pasku menu Środowisko wykonawcze > Uruchom sesję ponownie. Ponowne uruchomienie należy wykonać dopiero po zakończeniu wszystkich instalacji.
• Po ponownym uruchomieniu środowiska wykonawczego wykonaj tylko te polecenia . Nie musisz ponownie uruchamiać poprzednich.

4. Importowanie wszystkich bibliotek

Po zakończeniu instalacji przejdź do głównego wykonania.

Zalecenie: w tej i kolejnych sekcjach użyj notatnika Jupyter lub dodaj polecenia do skryptu Pythona.

Zaimportuj wymagane pakiety:

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

Uwaga: jeśli zamiast ai-edge-litert-nightly używasz stabilnego pakietu ai-edge-litert, przed wszystkimi instrukcjami importu (przy każdym uruchomieniu skryptu) musisz ustawić tę zmienną środowiskową: GOOGLE_TENSOR_BACKEND_ENABLED=1. Fragment kodu Pythona:

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

LUB (powłoka)

export GOOGLE_TENSOR_BACKEND_ENABLED=1

5. Kompilowanie modelu LiteRT

Ta sekcja obejmuje zaawansowane zastosowania, takie jak bezpośrednie kompilowanie modelu LiteRT (TFLite).

Kompilacja EdgeTPU z modelu TFLite

Ten krok wymaga modelu TFLite. Jeśli nie masz modelu TFLite, przekonwertuj go na format TFLite.

Pobieranie modelu TFLite

W tym przypadku używamy modelu MediaPipe MultiClass Segmentation.

Model TFLite jest dostępny na stronie MediaPipe segmentacja obrazu.

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)

Szybkie weryfikowanie modelu TfLite za pomocą interfejsu LiteRT Python API

W poniższym przykładzie zobaczysz zarówno obraz maski, jak i połączony wynik.

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

Konwertowanie na model LiteRT z kompilacją EdgeTPU AOT.

Do kompilowania modelu używamy interfejsów API z 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'
)

Po zakończeniu kompilacji użyj metody model.export, aby wyeksportować wszystkie modele na dysk.

Domyślnie modele są przechowywane w płaskiej strukturze w katalogu wyjściowym, a nazwa każdego modelu jest zakończona identyfikatorem backendu.

Na przykład:

Nazwa pliku modelu	Backend	SoC	Notatka
selfie_segmentation_fallback.tflite	Procesor/grafika	Nie dotyczy	Nie dotyczy
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. Eksportowanie i weryfikowanie na procesorze

Po zakończeniu kompilacji zweryfikuj model TFLite na procesorze. Zrób to za pomocą „modelu rezerwowego” wygenerowanego podczas kompilacji.

# 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. Eksportowanie modeli na potrzeby PODAI

Po zweryfikowaniu modeli kolejnym ważnym krokiem jest przygotowanie ich do wdrożenia. W tej sekcji dowiesz się, jak spakować skompilowane modele, aby przesłać je do Google Play, co umożliwi dostarczanie ich na urządzenia użytkowników za pomocą platformy Google Play On-Device AI (PODAI).

Moduł AiEdgeLiteRT AOT (Ahead-of-Time) udostępnia narzędzia ai_pack przeznaczone specjalnie do tego celu. Te narzędzia tworzą pakiet AI, który jest kluczowym zasobem danych. Pakiet AI łączy skompilowane modele z konfiguracjami kierowania na urządzenia, dzięki czemu odpowiednie modele i zasoby są dostarczane na odpowiednie urządzenia użytkowników. Jest to szczególnie ważne w przypadku kompilacji NPU (Neural Processing Unit), ponieważ zapewnia, że modele zoptymalizowane pod kątem konkretnego układu SoC (System-on-Chip) trafią tylko na urządzenia wyposażone w ten układ.

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

Sprawdzanie źródła pakietu 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. Konfigurowanie opcji zaawansowanych

Kompilacja NPU na konkretne urządzenie lub EdgeTPU

Domyślnie kompilacja LiteRT AOT kompiluje wszystkie zarejestrowane backendy. W przypadku lokalnego programowania możesz chcieć kompilować tylko na określonych urządzeniach, np. na telefonach deweloperskich. Aby to zrobić, podaj wyraźnie cele kompilacji.

Poniższy przykład kompiluje się do 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())

Flagi kompilacji dla Google Tensor

Dostosuj proces kompilacji za pomocą flag kompilacji. W tym przypadku używana jest ta flaga: google_tensor_truncation_type="half".

Podczas kompilowania modelu TFLite

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

9. Dalsze kroki

Gratulacje!

Twoje modele są gotowe do użycia przez PODAI.

Teraz przejdź do Android Studio, aby wykonać kolejne kroki. Szczegóły znajdziesz w przykładach segmentacji obrazu LiteRT.