Implementação do LiteRT no Google Tensor

1. Visão geral

O SDK do Google Tensor é usado para compilar modelos LiteRT para dispositivos Pixel. Os modelos compilados podem ser implantados em dispositivos Pixel para melhorar o desempenho da inferência de ML. Para usar o SDK, primeiro é preciso converter o modelo em um modelo LiteRT (tflite).

Este codelab é baseado no Colab geral no GitHub, Tutorial do Colab de compilação AOT do LiteRT:.

Objetivo

Saiba como usar o compilador LiteRT AOT (antecipado) para compilar um modelo de segmentação de selfie de um modelo do TFLite em um modelo do LiteRT otimizado e compilado para EdgeTPUs no dispositivo.

Este Colab também mostra as etapas para preparar modelos com o Play para IA no dispositivo (PODAI).

O PODAI oferece modelos personalizados para recursos de IA no dispositivo com mais eficiência. Ele simplifica o processo de lançamento, segmentação, controle de versões e download dos seus modelos de IA. Quando combinada com a compilação AOT LiteRT EdgeTPU, ela permite que os desenvolvedores entreguem modelos de ML compilados para vários dispositivos sem a necessidade de saber quais EdgeTPUs o smartphone do usuário final contém.

Modelos usados

Os modelos que estamos usando foram publicados originalmente no guia de segmentação de imagens do MediaPipe. Confira alguns detalhes sobre o modelo usado neste codelab:

• SelfieMulticlass: um modelo LiteRT que usa uma imagem de uma pessoa, localiza áreas como cabelo, pele e roupas e gera um mapa de segmentação de imagem para esses itens.

2. Primeiros passos

Siga estas etapas para acessar e começar a usar o SDK do Google Tensor:

1. Inscreva-se para ter acesso ao SDK do Google Tensor. Antes de continuar, aguarde um e-mail do Google com o link para download do plug-in do compilador.
2. Faça o download do plug-in do compilador (litert_plugin_compiler.tar.gz) e coloque-o em uma pasta de sua escolha.
3. Defina a variável de ambiente como o caminho do sistema local do arquivo baixado, GOOGLE_TENSOR_SDK_BETA.
   Execute este comando no terminal bash:

   export GOOGLE_TENSOR_SDK_BETA=/path/to/downloaded/compiler
   

   Ou execute isso no seu notebook do Colab:

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

4. Em seguida, execute este comando para instalar o pacote:

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

3. Instalar os pacotes necessários

Comece instalando os pacotes necessários, incluindo ai-edge-litert-nightly, que contém o compilador AOT da EdgeTPU, e outras bibliotecas usadas para conversão de modelos.

Use este pacote para instalar o back-end LiteRT para o Google Tensor: ai-edge-litert-sdk-google-tensor.

Depois de instalar os pacotes, reinicie a sessão e siga as etapas de instalação. Não repita a instalação.

Importante:se você planeja usar um notebook do Colab ou do Jupyter, adicione um ponto de exclamação ! antes de cada comando "pip install" e "pip uninstall". Faça isso ao copiar e colar o comando pip em uma célula do Colab.

Se você planeja realizar a configuração no seu sistema, recomendamos usar um ambiente virtual Python (venv) e executar esses comandos nele.

Desinstalar determinados pacotes

Antes disso, desinstale o tensorflow que vem com o ambiente de execução do Colab por padrão.

Observação:nesta etapa, você pode encontrar avisos sobre pacotes não instalados. Você pode ignorá-los com segurança.

pip uninstall -y tensorflow ai-edge-litert

Instalar todas as bibliotecas

Instalar o back-end LiteRT para o Google Tensor

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

Instalar os pacotes restantes

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

Importante:se você usa um notebook Colab ou Jupyter:

• Reinicie o ambiente de execução do Colab após a conclusão da instalação.
• Reinicie o ambiente de execução do Colab clicando em Ambiente de execução > Reiniciar sessão na barra de menus para que os pacotes instalados sejam refletidos corretamente. A reinicialização só deve ser feita depois que todas as instalações forem concluídas.
• Depois de reiniciar o ambiente de execução, execute apenas os comandos a seguir . Você não precisa executar os anteriores de novo.

4. Importar todas as bibliotecas

Prossiga para a execução principal depois que a instalação for concluída.

Recomendação:para esta e as próximas seções, use um notebook Jupyter ou adicione os comandos a um script Python.

Importe os pacotes necessários:

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

Observação:se você estiver usando o pacote de lançamento estável ai-edge-litert em vez de ai-edge-litert-nightly, defina esta variável de ambiente antes de todas as instruções de importação (para cada execução do script): GOOGLE_TENSOR_BACKEND_ENABLED=1. Snippet de código Python para o mesmo:

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

OR (shell)

export GOOGLE_TENSOR_BACKEND_ENABLED=1

5. Compilar um modelo LiteRT

Esta seção aborda usos avançados, como a compilação direta de um modelo LiteRT (TFLite).

Compilação do EdgeTPU com base no modelo do TFLite

Esta etapa exige um modelo do TFLite. Se você não tiver um modelo do TFLite, converta o modelo para esse formato.

Receber o modelo do TFLite

Usamos o modelo MediaPipe MultiClass Segmentation para esse caso de uso.

O modelo do TFLite está disponível na página Segmentação de imagens do 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)

Verificar rapidamente o modelo do TfLite usando a API Python LiteRT

No exemplo a seguir, você verá a imagem da máscara e o 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()

Converter para o modelo LiteRT com compilação AOT do EdgeTPU.

Usamos as APIs de ai_edge_litert.aot para compilar o 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'
)

Quando a compilação terminar, use o método model.export para exportar todos os modelos para o disco.

Por padrão, os modelos são armazenados em uma estrutura simples no diretório de saída, com o nome de cada modelo sufixado com o ID do back-end.

Exemplo:

Nome do arquivo do modelo	Back-end	SoC	Observação
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. Exportar e validar na CPU

Depois que a compilação for concluída, verifique o modelo do TFLite na CPU. Faça isso usando o "modelo de substituição" gerado durante a compilação.

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

Com os modelos verificados, a próxima etapa essencial é prepará-los para a implantação. Esta seção detalha como empacotar seus modelos compilados para fazer upload no Google Play, permitindo a entrega aos dispositivos dos usuários pelo framework de IA no dispositivo do Google Play (PODAI, na sigla em inglês).

O módulo AOT (Ahead-of-Time) AiEdgeLiteRT oferece utilitários ai_pack especificamente para essa finalidade. Esses utilitários criam um pacote de IA, que é um recurso de dados crucial. Um pacote de IA agrupa seus modelos compilados com configurações de segmentação por dispositivo, garantindo que os modelos e recursos corretos sejam entregues aos dispositivos de usuário adequados. Isso é particularmente importante para compilações de NPU (unidade de processamento neural), porque garante que os modelos otimizados para um System-on-Chip (SoC) específico cheguem apenas aos dispositivos equipados com esse 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
)

Inspeção da origem do pacote 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. Configure opções avançadas

Compilação de NPU para dispositivo específico ou EdgeTPU

Por padrão, a compilação AOT do LiteRT é feita para todos os backends registrados. Para desenvolvimento local, talvez você queira compilar apenas para dispositivos específicos, como smartphones de desenvolvimento. Para isso, forneça os destinos de compilação explicitamente.

O exemplo a seguir é compilado para o 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())

Flags de compilação para o Google Tensor

Personalize o processo de compilação com flags de compilação. Aqui, esta flag é usada: google_tensor_truncation_type="half"

Ao compilar um modelo do TFLite

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

9. Próximas etapas

Parabéns!

Seus modelos estão prontos para serem usados pela PODAI.

Agora, siga para o Android Studio para as próximas etapas. Consulte Exemplos de segmentação de imagens do LiteRT para mais detalhes.