Реализация LiteRT в Google Tensor

1. Обзор

SDK Google Tensor используется для компиляции моделей LiteRT для устройств Pixel. Скомпилированные модели можно развернуть на устройствах Pixel для повышения производительности машинного обучения. Для использования SDK необходимо сначала преобразовать вашу модель в модель LiteRT ( tflite ).

Данный практический пример основан на общем проекте Colab на GitHub, LiteRT AOT Compilation Tutorial Colab: .

Цель

Узнайте, как использовать компилятор LiteRT AOT (ahead of time) для компиляции модели сегментации селфи из модели TFLite в модель LiteRT, оптимизированную и скомпилированную для встроенных EdgeTPU.

В этом руководстве также описаны шаги по подготовке моделей с помощью Play for On-device AI ( PODAI ).

PODAI обеспечивает более эффективную доставку пользовательских моделей для функций искусственного интеллекта на устройстве. Он упрощает процесс запуска, выбора целевой платформы, версионирования и загрузки моделей ИИ. В сочетании с компиляцией LiteRT EdgeTPU AOT он позволяет разработчикам предоставлять скомпилированные модели машинного обучения для различных устройств без необходимости знать, какие EdgeTPU установлены на телефоне конечного пользователя.

Используемые модели

Используемые нами модели первоначально были опубликованы в руководстве по сегментации изображений MediaPipe. Вот некоторые подробности о модели, используемой в этом практическом занятии:

• SelfieMulticlass : модель LiteRT, которая принимает изображение человека, определяет области, такие как волосы, кожа и одежда, и выдает карту сегментации изображения для этих элементов.

2. Начните

Выполните следующие шаги, чтобы получить доступ к Google Tensor SDK и начать работу с ним:

1. Зарегистрируйтесь , чтобы получить доступ к Google Tensor SDK. Прежде чем продолжить, вам необходимо дождаться электронного письма от Google со ссылкой для скачивания плагина компилятора.
2. Загрузите плагин компилятора ( litert_plugin_compiler.tar.gz ) и поместите его в папку по вашему выбору.
3. Установите переменную среды, указав локальный системный путь к загруженному файлу: GOOGLE_TENSOR_SDK_BETA .
   Вы можете выполнить эту команду в терминале bash:

   export GOOGLE_TENSOR_SDK_BETA=/path/to/downloaded/compiler
   

   Или вы можете запустить это в своем блокноте Colab:

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

4. Затем выполните следующую команду для установки пакета:

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

3. Установите необходимые пакеты.

Для начала установите необходимые пакеты, включая ai-edge-litert-nightly , который содержит компилятор EdgeTPU AOT, и другие библиотеки, используемые для преобразования моделей.

Для установки бэкенда LiteRT для Google Tensor используйте этот пакет: ai-edge-litert-sdk-google-tensor .

После установки пакетов перезапустите сессию и продолжите выполнение шагов установки. Не повторяйте установку.

Важно: Если вы планируете использовать блокнот Colab или Jupyter, добавьте восклицательный знак ! перед каждой командой pip install и pip uninstall. Делайте это при копировании и вставке команды pip в ячейку Colab.

Если вы планируете выполнить настройку в своей системе, мы рекомендуем использовать виртуальное окружение Python (venv) и запускать эти команды внутри виртуального окружения.

Удалите определенные пакеты.

Перед этим удалите TensorFlow, который по умолчанию входит в состав среды выполнения Colab.

Примечание: На этом этапе вы можете столкнуться с предупреждениями о том, что пакеты не установлены. Вы можете смело их игнорировать.

pip uninstall -y tensorflow ai-edge-litert

Установите все библиотеки.

Установите бэкенд LiteRT для Google Tensor.

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

Установите оставшиеся пакеты.

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

Важно: Если вы используете блокнот Colab или Jupyter:

• После завершения установки перезапустите среду выполнения Colab.
• Для корректного отображения установленных пакетов перезапустите среду выполнения Colab, выбрав в строке меню «Среда выполнения» > «Перезапустить сессию» . Перезапуск следует выполнять только после завершения всех установок.
• После перезапуска среды выполнения выполните только следующие команды . Повторно запускать предыдущие команды не нужно.

4. Импортируйте все библиотеки.

После завершения установки перейдите к основному выполнению.

Рекомендация: Для этого и последующих разделов используйте блокнот Jupyter или добавьте команды в скрипт Python.

Импортируйте необходимые пакеты:

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

Примечание: Если вы используете стабильный пакет ai-edge-litert вместо ai-edge-litert-nightly , то вам потребуется установить эту переменную окружения перед всеми операторами импорта (при каждом запуске скрипта): GOOGLE_TENSOR_BACKEND_ENABLED=1 . Пример кода на Python:

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

ИЛИ (оболочка)

export GOOGLE_TENSOR_BACKEND_ENABLED=1

5. Скомпилируйте модель LiteRT.

В этом разделе рассматриваются расширенные возможности, такие как прямая компиляция модели LiteRT (TFLite).

Компиляция EdgeTPU из модели TFLite

Для этого шага необходима модель в формате TFLite. Если у вас нет модели в формате TFLite, преобразуйте её в этот формат.

Приобретите модель TFLite.

Для решения этой задачи мы используем модель многоклассовой сегментации MediaPipe .

Модель TFLite доступна на странице сегментации изображений 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)

Быстрая проверка модели TfLite с помощью Python API LiteRT.

В следующем примере вы увидите как изображение-маску, так и результат смешивания.

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

Преобразование в модель LiteRT с компиляцией EdgeTPU AOT.

Для компиляции модели мы используем API из 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'
)

После завершения компиляции используйте метод model.export для экспорта всех моделей на диск.

По умолчанию модели хранятся в выходном каталоге в плоской структуре, при этом к имени каждой модели добавляется идентификатор бэкэнда.

Например:

Имя файла модели	Бэкенд	SoC	Примечание
selfie_segmentation_fallback.tflite	ЦП/ГП	Н/Д	Н/Д
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. Экспорт и проверка на ЦП.

После завершения компиляции проверьте модель TFLite на процессоре. Сделайте это, используя "резервную модель", сгенерированную во время компиляции.

# 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. Модели экспорта для PODAI

После проверки ваших моделей следующим важным шагом является их подготовка к развертыванию. В этом разделе подробно описано, как упаковать скомпилированные модели для загрузки в Google Play, что позволит доставлять их на устройства пользователей через фреймворк Google Play On-Device AI (PODAI).

Модуль AiEdgeLiteRT AOT (Ahead-of-Time) предоставляет утилиты ai_pack специально предназначенные для этой цели. Эти утилиты создают пакет AI Pack , который является важнейшим ресурсом данных. Пакет AI Pack объединяет скомпилированные модели с конфигурациями, ориентированными на конкретные устройства, обеспечивая доставку правильных моделей и ресурсов на соответствующие пользовательские устройства. Это особенно важно для компиляции на нейронных процессорах (NPU), поскольку гарантирует, что модели, оптимизированные для конкретной системы на кристалле (SoC), будут достигать только устройств, оснащенных этой 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
)

Проверка исходного кода пакета ИИ

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. Настройка дополнительных параметров

Компиляция NPU для конкретного устройства или EdgeTPU

По умолчанию компиляция LiteRT AOT выполняется для всех зарегистрированных бэкендов. Для локальной разработки может потребоваться компиляция только для определенных устройств, например, для мобильных телефонов разработчиков. Этого можно добиться, явно указав цели компиляции.

Следующий пример компилируется в 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())

Флаги компиляции для Google Tensor

Настройте процесс компиляции с помощью флагов компиляции. Здесь используется следующий флаг: google_tensor_truncation_type="half"

При компиляции модели TFLite

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

9. Дальнейшие шаги

Поздравляем!

Ваши модели готовы к использованию системой PODAI!

Теперь перейдите в Android Studio и выполните следующие шаги; подробности см. в примерах сегментации изображений LiteRT .