Implémenter LiteRT sur Google Tensor

1. Présentation

Le SDK Google Tensor est utilisé pour compiler les modèles LiteRT pour les appareils Pixel. Les modèles compilés peuvent être déployés sur des appareils Pixel pour améliorer les performances d'inférence du ML. Pour utiliser le SDK, vous devez d'abord convertir votre modèle en modèle LiteRT (tflite).

Cet atelier de programmation s'appuie sur le notebook Colab général sur GitHub, Tutoriel Colab sur la compilation AOT LiteRT : .

Objectif

Découvrez comment utiliser le compilateur LiteRT AOT (Ahead-Of-Time) pour compiler un modèle de segmentation de selfies à partir d'un modèle TFLite dans un modèle LiteRT optimisé et compilé pour les EdgeTPU sur l'appareil.

Ce notebook Colab vous explique également comment préparer des modèles avec Play for On-device AI (PODAI).

PODAI fournit des modèles personnalisés pour les fonctionnalités d'IA sur l'appareil de manière plus efficace. Il simplifie le processus de lancement, de ciblage, de gestion des versions et de téléchargement de vos modèles d'IA. Combinée à la compilation AOT LiteRT EdgeTPU, elle permet aux développeurs de fournir des modèles ML compilés pour différents appareils sans avoir à connaître les EdgeTPU contenus dans le téléphone de l'utilisateur final.

Modèles utilisés

Les modèles que nous utilisons sont initialement publiés dans le guide de segmentation d'images MediaPipe. Voici quelques détails concernant le modèle utilisé dans cet atelier de programmation :

• SelfieMulticlass : modèle LiteRT qui prend une image d'une personne, localise des zones telles que les cheveux, la peau et les vêtements, et génère une carte de segmentation d'image pour ces éléments.

2. Premiers pas

Pour accéder au SDK Google Tensor et commencer à l'utiliser, procédez comme suit :

1. Inscrivez-vous pour accéder au SDK Google Tensor. Avant de continuer, vous devez attendre de recevoir un e-mail de Google contenant le lien de téléchargement du plug-in du compilateur.
2. Téléchargez le plug-in du compilateur (litert_plugin_compiler.tar.gz) et placez-le dans le dossier de votre choix.
3. Définissez la variable d'environnement sur le chemin d'accès au système local du fichier téléchargé, GOOGLE_TENSOR_SDK_BETA.
    Vous pouvez exécuter cette commande dans votre terminal Bash :

   export GOOGLE_TENSOR_SDK_BETA=/path/to/downloaded/compiler
   

   Vous pouvez également exécuter cette commande dans votre notebook Colab :

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

4. Exécutez ensuite cette commande pour installer le package :

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

3. Installer les packages requis

Commencez par installer les packages requis, y compris ai-edge-litert-nightly, qui contient le compilateur AOT EdgeTPU, ainsi que les autres bibliothèques que vous utilisez pour la conversion de modèles.

Utilisez ce package pour installer le backend LiteRT pour Google Tensor : ai-edge-litert-sdk-google-tensor.

Une fois les packages installés, redémarrez la session et reprenez les étapes d'installation. Ne réessayez pas l'installation.

Important : Si vous prévoyez d'utiliser un notebook Colab ou Jupyter, ajoutez un point d'exclamation ! avant chaque commande pip install et pip uninstall. Faites-le lorsque vous copiez et collez la commande pip dans une cellule Colab.

Si vous prévoyez d'effectuer la configuration sur votre système, nous vous recommandons d'utiliser un environnement virtuel Python (venv) et d'exécuter ces commandes dans l'environnement virtuel.

Désinstaller certains packages

Avant cela, désinstallez TensorFlow fourni avec l'environnement d'exécution Colab par défaut.

Remarque : Lors de cette étape, vous pouvez rencontrer des avertissements concernant des packages non installés. Vous pouvez les ignorer sans problème.

pip uninstall -y tensorflow ai-edge-litert

Installer toutes les bibliothèques

Installer le backend LiteRT pour Google Tensor

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

Installer les packages restants

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

Important : Si vous utilisez un notebook Colab ou Jupyter :

• Redémarrez l'environnement d'exécution Colab une fois l'installation terminée.
• Redémarrez l'environnement d'exécution de Colab en cliquant sur Environnement d'exécution > Redémarrer la session dans la barre de menu pour que les packages installés soient correctement reflétés. Le redémarrage ne doit être effectué qu'une fois toutes les installations terminées.
• Après avoir redémarré le runtime, exécutez uniquement les commandes suivantes . Vous n'avez pas besoin de réexécuter les commandes précédentes.

4. Importer toutes les bibliothèques

Passez à l'exécution principale une fois l'installation terminée.

Recommandation : Pour cette section et les suivantes, utilisez un notebook Jupyter ou ajoutez les commandes à un script Python.

Importez les packages requis :

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

Remarque : Si vous utilisez le package de version stable ai-edge-litert au lieu de ai-edge-litert-nightly, vous devrez définir cette variable d'environnement avant toutes les instructions d'importation (pour chaque exécution du script) : GOOGLE_TENSOR_BACKEND_ENABLED=1. Extrait de code Python correspondant :

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

OR (shell)

export GOOGLE_TENSOR_BACKEND_ENABLED=1

5. Compiler un modèle LiteRT

Cette section aborde les utilisations avancées, comme la compilation directe d'un modèle LiteRT (TFLite).

Compilation EdgeTPU à partir du modèle TFLite

Cette étape nécessite un modèle TFLite. Si vous n'avez pas de modèle TFLite, convertissez votre modèle au format TFLite.

Obtenir le modèle TFLite

Pour ce cas d'utilisation, nous utilisons le modèle MediaPipe MultiClass Segmentation.

Le modèle TFLite est disponible sur la page Segmentation d'image 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)

Vérifier rapidement le modèle TfLite à l'aide de l'API Python LiteRT

Dans l'exemple suivant, vous verrez à la fois l'image du masque et le résultat du mélange.

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

Convertissez-le en modèle LiteRT, avec la compilation AOT EdgeTPU.

Nous utilisons les API de ai_edge_litert.aot pour compiler le modèle.

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

Une fois la compilation terminée, utilisez la méthode model.export pour exporter tous les modèles sur le disque.

Par défaut, les modèles sont stockés dans une structure plate du répertoire de sortie, chaque nom de modèle étant suffixé par l'ID du backend.

Exemple :

Nom de fichier du modèle	Backend	SoC (System on Chip)	Remarque
selfie_segmentation_fallback.tflite	Processeur/GPU	N/A	N/A
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. Exporter et valider sur le CPU

Une fois la compilation terminée, vérifiez le modèle TFLite sur le processeur. Pour ce faire, utilisez le "modèle de secours" généré lors de la compilation.

# 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. Exporter des modèles pour PODAI

Une fois vos modèles vérifiés, l'étape essentielle suivante consiste à les préparer au déploiement. Cette section explique comment empaqueter vos modèles compilés pour les importer sur Google Play, ce qui permet de les distribuer aux appareils des utilisateurs via le framework Google Play On-Device AI (PODAI).

Le module AOT (Ahead-of-Time) AiEdgeLiteRT fournit des utilitaires ai_pack spécialement conçus à cet effet. Ces utilitaires créent un pack d'IA, qui est un élément de données essentiel. Un pack d'IA regroupe vos modèles compilés avec des configurations de ciblage des appareils, ce qui garantit que les modèles et les composants appropriés sont fournis aux appareils des utilisateurs concernés. Cela est particulièrement important pour les compilations NPU (Neural Processing Unit), car cela garantit que les modèles optimisés pour un System-on-Chip (SoC) spécifique n'atteignent que les appareils équipés de ce 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
)

Inspecter la source du pack IA

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. Configurez les options avancées.

Compilation de NPU pour un appareil ou un EdgeTPU spécifique

Par défaut, la compilation LiteRT AOT compile tous les backends enregistrés. Pour le développement local, vous ne souhaiterez peut-être compiler que pour des appareils spécifiques, tels que des téléphones de développement. Pour ce faire, indiquez explicitement les cibles de compilation.

L'exemple suivant est compilé sur 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())

Options de compilation pour Google Tensor

Personnalisez le processus de compilation à l'aide de flags de compilation. Ici, cette option est utilisée : google_tensor_truncation_type="half"

Lors de la compilation d'un modèle TFLite

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

9. Étapes suivantes

Félicitations !

Vos modèles sont prêts à être utilisés par PODAI.

Passez maintenant à Android Studio pour les étapes suivantes. Pour en savoir plus, consultez Exemples de segmentation d'images LiteRT.