Implementazione di LiteRT su Google Tensor

1. Panoramica

L'SDK Google Tensor viene utilizzato per compilare i modelli LiteRT per i dispositivi Pixel. I modelli compilati possono essere implementati sui dispositivi Pixel per migliorare le prestazioni di inferenza ML. Per utilizzare l'SDK, devi prima convertire il modello in un modello LiteRT (tflite).

Questo codelab si basa sul colab generale su GitHub, LiteRT AOT Compilation Tutorial Colab:.

Obiettivo

Scopri come utilizzare il compilatore LiteRT AOT (ahead of time) per compilare un modello di segmentazione dei selfie da un modello TFLite in un modello LiteRT ottimizzato e compilato per EdgeTPU sul dispositivo.

Questo colab ti guida anche nei passaggi per preparare i modelli con Play for On-device AI (PODAI).

PODAI offre modelli personalizzati per le funzionalità di AI sul dispositivo in modo più efficiente. Semplifica il processo di lancio, targeting, controllo delle versioni e download dei modelli di AI. Se combinato con la compilazione AOT di LiteRT EdgeTPU, consente agli sviluppatori di fornire modelli ML compilati per vari dispositivi senza dover sapere quali EdgeTPU contiene lo smartphone dell'utente finale.

Modelli utilizzati

I modelli che utilizziamo sono stati pubblicati originariamente nella guida alla segmentazione delle immagini di MediaPipe Image segmentation guide. Ecco alcuni dettagli sul modello utilizzato in questo codelab:

• SelfieMulticlass: un modello LiteRT che prende un'immagine di una persona, individua aree come capelli, pelle e vestiti e restituisce una mappa di segmentazione delle immagini per questi elementi.

2. Inizia

Per accedere all'SDK Google Tensor e iniziare a utilizzarlo, segui questi passaggi:

1. Registrati per accedere all'SDK Google Tensor. Prima di procedere, devi attendere un'email da Google contenente il link per il download del plug-in del compilatore.
2. Scarica il plug-in del compilatore (litert_plugin_compiler.tar.gz) e inseriscilo in una cartella a tua scelta.
3. Imposta la variabile di ambiente sul percorso del sistema locale del file scaricato, GOOGLE_TENSOR_SDK_BETA.
   Puoi eseguire questo comando nel terminale bash:

   export GOOGLE_TENSOR_SDK_BETA=/path/to/downloaded/compiler
   

   In alternativa, puoi eseguirlo nel notebook Colab:

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

4. Quindi esegui questo comando per installare il pacchetto:

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

3. Installare i pacchetti richiesti

Inizia installando i pacchetti richiesti, tra cui ai-edge-litert-nightly, che contiene il compilatore AOT EdgeTPU, e altre librerie utilizzate per la conversione dei modelli.

Utilizza questo pacchetto per installare il backend LiteRT per Google Tensor: ai-edge-litert-sdk-google-tensor.

Dopo aver installato i pacchetti, riavvia la sessione e procedi con i passaggi di installazione. Non ripetere l'installazione.

Importante: se prevedi di utilizzare un notebook di Colab o un notebook Jupyter, aggiungi un punto esclamativo ! prima di ogni comando pip install e pip uninstall. Esegui questa operazione quando copi e incolli il comando pip in una cella colab.

Se prevedi di eseguire la configurazione sul tuo sistema, ti consigliamo di utilizzare un ambiente virtuale Python (venv) ed eseguire questi comandi all'interno dell'ambiente virtuale.

Disinstallare determinati pacchetti

Prima di tutto, disinstalla tensorflow fornito con il runtime colab per impostazione predefinita.

Nota: in questo passaggio potresti visualizzare avvisi relativi a pacchetti non installati. Puoi ignorarli in sicurezza.

pip uninstall -y tensorflow ai-edge-litert

Installare tutte le librerie

Installare il backend LiteRT per Google Tensor

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

Installare i pacchetti rimanenti

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

Importante: se utilizzi un notebook Colab o Jupyter:

• Riavvia il runtime Colab al termine dell'installazione.
• Riavvia il runtime di Colab facendo clic su Runtime > Riavvia sessione dalla barra dei menu per visualizzare correttamente i pacchetti installati. Il riavvio deve essere eseguito solo al termine di tutte le installazioni.
• Dopo aver riavviato il runtime, esegui solo i comandi che seguono . Non devi eseguire di nuovo quelli precedenti.

4. Importare tutte le librerie

Procedi all'esecuzione principale al termine dell'installazione.

Consiglio: per questa e le sezioni seguenti, utilizza un notebook Jupyter o aggiungi i comandi a uno script Python.

Importa i pacchetti richiesti:

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

Nota: se utilizzi il pacchetto di release stabile ai-edge-litert anziché ai-edge-litert-nightly, dovrai impostare questa variabile di ambiente prima di tutte le istruzioni di importazione (per ogni esecuzione dello script): GOOGLE_TENSOR_BACKEND_ENABLED=1. Snippet di codice Python per lo stesso:

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

OPPURE (shell)

export GOOGLE_TENSOR_BACKEND_ENABLED=1

5. Compilare un modello LiteRT

Questa sezione illustra gli utilizzi avanzati, ad esempio la compilazione diretta di un modello LiteRT (TFLite).

Compilazione EdgeTPU dal modello TFLite

Questo passaggio richiede un modello TFLite. Se non hai un modello TFLite, converti il modello nel formato TFLite.

Recuperare il modello TFLite

Per questo caso d'uso utilizziamo il modello di segmentazione multiclasse MediaPipe.

Il modello TFLite è disponibile nella pagina Segmentazione delle immagini di 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)

Verificare rapidamente il modello TfLite utilizzando l'API Python LiteRT

Nell'esempio seguente vedrai sia l'immagine della maschera sia il risultato combinato.

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

Convertire in modello LiteRT con compilazione AOT EdgeTPU.

Utilizziamo le API di ai_edge_litert.aot per compilare il modello.

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

Al termine della compilazione, utilizza il metodo model.export per esportare tutti i modelli su disco.

Per impostazione predefinita, i modelli vengono archiviati in una struttura piatta nella directory di output, con il nome di ogni modello con il suffisso dell'ID del backend.

Ad esempio:

Nome file del modello	Backend	SoC	Nota
selfie_segmentation_fallback.tflite	CPU/GPU	N/D	N/D
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. Esportare e convalidare sulla CPU

Al termine della compilazione, verifica il modello TFLite sulla CPU. A questo scopo, utilizza il "modello di fallback" generato durante la compilazione.

# 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. Esportare i modelli per PODAI

Dopo aver verificato i modelli, il passaggio essenziale successivo è prepararli per il deployment. Questa sezione descrive in dettaglio come creare pacchetti dei modelli compilati per il caricamento su Google Play, consentendo la distribuzione ai dispositivi degli utenti tramite il framework Google Play On-Device AI (PODAI).

Il modulo AiEdgeLiteRT AOT (Ahead-of-Time) fornisce le utilità ai_pack appositamente per questo scopo. Queste utilità creano un pacchetto AI, che è un asset di dati fondamentale. Un pacchetto AI raggruppa i modelli compilati con le configurazioni di targeting dei dispositivi, assicurando che i modelli e gli asset corretti vengano distribuiti ai dispositivi utente appropriati. Questo è particolarmente importante per le compilazioni NPU (Neural Processing Unit), in quanto garantisce che i modelli ottimizzati per un System-on-Chip (SoC) specifico raggiungano solo i dispositivi dotati di quel 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
)

Esaminare l'origine del pacchetto 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. Configurare le opzioni avanzate

Compilazione NPU per un dispositivo o un EdgeTPU specifico

Per impostazione predefinita, la compilazione AOT di LiteRT viene eseguita su tutti i backend registrati. Per lo sviluppo locale, potresti voler compilare solo per dispositivi specifici, ad esempio smartphone di sviluppo. Per farlo, fornisci esplicitamente le destinazioni di compilazione.

L'esempio seguente compila per 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 di compilazione per Google Tensor

Personalizza il processo di compilazione tramite i flag di compilazione. Qui viene utilizzato questo flag: google_tensor_truncation_type="half"

Quando compili un modello TFLite

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

9. Passaggi successivi

Congratulazioni!

I tuoi modelli sono pronti per essere utilizzati da PODAI.

Ora passa ad Android Studio per i passaggi successivi. Per i dettagli, consulta gli esempi di segmentazione delle immagini LiteRT.