Implementación de LiteRT en Google Tensor

1. Descripción general

El SDK de Google Tensor se usa para compilar modelos de LiteRT para dispositivos Pixel. Los modelos compilados se pueden implementar en dispositivos Pixel para mejorar el rendimiento de la inferencia de AA. Para usar el SDK, primero debes convertir tu modelo en un modelo de LiteRT (tflite).

Este codelab se basa en el colab general en GitHub, LiteRT AOT Compilation Tutorial Colab:.

Objetivo

Aprende a usar el compilador LiteRT AOT (anticipado) para compilar un modelo de segmentación de selfies de un modelo de TFLite en un modelo de LiteRT optimizado y compilado para EdgeTPUs integradas en el dispositivo.

Este colab también te guía por los pasos para preparar modelos con Play para IA integrada en el dispositivo (PODAI).

PODAI ofrece modelos personalizados para funciones de IA integradas en el dispositivo de manera más eficiente. Simplifica el proceso de lanzamiento, segmentación, control de versiones y descarga de tus modelos de IA. Cuando se combina con la compilación AOT de LiteRT EdgeTPU, permite a los desarrolladores entregar modelos de AA compilados para varios dispositivos sin necesidad de saber qué EdgeTPUs contiene el teléfono del usuario final.

Modelos usados

Los modelos que usamos se publicaron originalmente en la guía de segmentación de imágenes de MediaPipe. Estos son algunos detalles sobre el modelo que se usa en este codelab:

• SelfieMulticlass: Es un modelo de LiteRT que toma una imagen de una persona, ubica áreas como el cabello, la piel y la ropa, y muestra un mapa de segmentación de imágenes para estos elementos.

2. Comenzar

Sigue estos pasos para obtener acceso y comenzar a usar el SDK de Google Tensor:

1. Regístrate para obtener acceso al SDK de Google Tensor. Antes de continuar, debes esperar un correo electrónico de Google que contenga el vínculo de descarga del complemento del compilador.
2. Descarga el complemento del compilador (litert_plugin_compiler.tar.gz) y colócalo en la carpeta que elijas.
3. Establece la variable de entorno en la ruta de acceso del sistema local del archivo descargado, GOOGLE_TENSOR_SDK_BETA.
   Puedes ejecutar este comando en tu terminal de bash:

   export GOOGLE_TENSOR_SDK_BETA=/path/to/downloaded/compiler
   

   O puedes ejecutarlo en tu notebook de Colab:

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

4. Luego, ejecuta este comando para instalar el paquete:

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

3. Instala los paquetes obligatorios

Comienza por instalar los paquetes obligatorios, incluido ai-edge-litert-nightly, que contiene el compilador AOT de EdgeTPU y otras bibliotecas que usas para la conversión de modelos.

Usa este paquete para instalar el backend de LiteRT para Google Tensor: ai-edge-litert-sdk-google-tensor.

Después de instalar los paquetes, reinicia la sesión y continúa con los pasos de instalación. No repitas la instalación.

Importante: Si planeas usar un notebook de Colab o un notebook de Jupyter, agrega un signo de exclamación ! antes de cada comando pip install y pip uninstall. Haz esto cuando copies y pegues el comando pip en una celda de Colab.

Si planeas realizar la configuración en tu sistema, te recomendamos que uses un entorno virtual de Python (venv) y ejecutes estos comandos dentro del entorno virtual.

Desinstala ciertos paquetes

Antes de eso, desinstala el tensorflow que viene con el entorno de ejecución de Colab de forma predeterminada.

Nota: En este paso, es posible que veas advertencias sobre paquetes que no están instalados. Puedes ignorarlos de forma segura.

pip uninstall -y tensorflow ai-edge-litert

Instala todas las bibliotecas

Instala el backend de LiteRT para Google Tensor

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

Instala los paquetes restantes

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

Importante: Si usas un notebook de Colab o Jupyter, haz lo siguiente:

• Reinicia el entorno de ejecución de Colab después de que se complete la instalación.
• Para reiniciar el entorno de ejecución de Colab, haz clic en Entorno de ejecución > Reiniciar sesión en la barra de menú para que los paquetes instalados se reflejen correctamente. El reinicio solo debe realizarse después de que se completen todas las instalaciones.
• Después de reiniciar el entorno de ejecución, ejecuta solo los comandos que siguen . No es necesario que vuelvas a ejecutar los anteriores.

4. Importa todas las bibliotecas

Continúa con la ejecución principal después de que se complete la instalación.

Recomendación: Para esta y las siguientes secciones, usa un notebook de Jupyter o agrega los comandos a una secuencia de comandos de Python.

Importa los paquetes obligatorios:

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: Si usas el paquete de versión estable ai-edge-litert en lugar de ai-edge-litert-nightly, deberás establecer esta variable de entorno antes de todas las instrucciones de importación (para cada ejecución de la secuencia de comandos): GOOGLE_TENSOR_BACKEND_ENABLED=1. Fragmento de código de Python para lo mismo:

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

O (shell)

export GOOGLE_TENSOR_BACKEND_ENABLED=1

5. Compila un modelo de LiteRT

En esta sección, se abarcan usos avanzados, como la compilación directa de un modelo de LiteRT (TFLite).

Compilación de EdgeTPU desde el modelo de TFLite

Este paso requiere un modelo de TFLite. Si no tienes un modelo de TFLite, convierte tu modelo al formato de TFLite.

Obtén el modelo de TFLite

Usamos el modelo de segmentación de varias clases de MediaPipe para este caso de uso.

El modelo de TFLite está disponible en la página de segmentación de imágenes de 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)

Verifica rápidamente el modelo de TfLite con la API de Python de LiteRT

En el siguiente ejemplo, verás la imagen de la máscara y el resultado combinado.

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

Convierte al modelo de LiteRT con la compilación AOT de EdgeTPU.

Usamos las APIs de ai_edge_litert.aot para compilar el modelo.

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

Cuando finalice la compilación, usa el método model.export para exportar todos los modelos al disco.

De forma predeterminada, los modelos se almacenan en una estructura plana en el directorio de salida, y cada nombre de modelo tiene el sufijo del ID de backend.

Por ejemplo:

Nombre de archivo del modelo	Backend	SoC	Nota
selfie_segmentation_fallback.tflite	CPU/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. Exporta y valida en la CPU

Una vez que se complete la compilación, verifica el modelo de TFLite en la CPU. Para ello, usa el "modelo de resguardo" generado durante la compilación.

# 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. Exporta modelos para PODAI

Una vez que se verifican los modelos, el siguiente paso esencial es prepararlos para la implementación. En esta sección, se explica cómo empaquetar los modelos compilados para subirlos a Google Play, lo que permite la entrega a los dispositivos de los usuarios a través del framework de IA integrada en el dispositivo de Google Play (PODAI).

El módulo AOT (Ahead-of-Time) de AiEdgeLiteRT proporciona utilidades ai_pack específicamente para este propósito. Estas utilidades crean un paquete de IA, que es un recurso de datos crucial. Un paquete de IA agrupa tus modelos compilados con configuraciones de segmentación de dispositivos, lo que garantiza que los modelos y recursos correctos se entreguen a los dispositivos de usuario adecuados. Esto es especialmente importante para las compilaciones de NPU (unidad de procesamiento neuronal), ya que garantiza que los modelos optimizados para un sistema en chip (SoC) específico lleguen solo a los dispositivos equipados con ese 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
)

Inspecciona la fuente del paquete de 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. Configura opciones avanzadas.

Compilación de NPU para un dispositivo o EdgeTPU específicos

De forma predeterminada, la compilación AOT de LiteRT se compila en todos los backends registrados. Para el desarrollo local, es posible que solo quieras compilar para dispositivos específicos, como teléfonos de desarrollo. Para lograr esto, proporciona los objetivos de compilación de forma explícita.

En el siguiente ejemplo, se compila en 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())

Marcas de compilación para Google Tensor

Personaliza el proceso de compilación a través de marcas de compilación. Aquí, se usa esta marca: google_tensor_truncation_type="half"

Cuando se compila un modelo de TFLite

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

9. Próximos pasos

Felicitaciones

Tus modelos están listos para que PODAI los consuma.

Ahora, pasa a Android Studio para los siguientes pasos. Consulta los ejemplos de segmentación de imágenes de LiteRT para obtener más detalles.