在 Google Tensor 上实现 LiteRT

1. 概览

Google Tensor SDK 用于为 Pixel 设备编译 LiteRT 模型。编译后的模型可以部署在 Pixel 设备上，以提升机器学习推断性能。如需使用该 SDK，您必须先将模型转换为 LiteRT (tflite) 模型。

此 Codelab 基于 GitHub 上的通用 Colab LiteRT AOT 编译教程 Colab:在 Colab 中打开。

目标

了解如何使用 LiteRT AOT （提前）编译器将 自拍分割模型 从 TFLite 模型编译为针对设备端 EdgeTPU 进行了优化和编译的 LiteRT 模型。

此 Colab 还将引导您完成使用 Play for On-device AI (PODAI) 准备模型的步骤。

PODAI 可以更高效地为设备端 AI 功能提供自定义模型。它简化了 AI 模型的启动、定位、版本控制和下载流程。与 LiteRT EdgeTPU AOT 编译相结合，开发者可以为各种设备提供已编译的机器学习模型，而无需了解最终用户的手机包含哪些 EdgeTPU。

使用的模型

我们使用的模型最初发布在 MediaPipe 图像分割指南中。以下是有关此 Codelab 中使用的模型的一些详细信息：

  • 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 编译器）以及您用于模型转换的其他库。

使用此软件包为 Google Tensor 安装 LiteRT 后端：ai-edge-litert-sdk-google-tensor。

安装软件包后，重启会话并从安装步骤继续。请勿重复安装。

如果您打算在系统上执行设置，建议您使用 Python 虚拟环境 (venv) 并在虚拟环境中运行这些命令。

卸载某些软件包

在此之前，请先卸载 Colab 运行时默认附带的 TensorFlow。

pip uninstall -y tensorflow ai-edge-litert

安装所有库

为 Google Tensor 安装 LiteRT 后端

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

安装剩余的软件包

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

4. 导入所有库

安装完成后，继续执行主要执行操作。

导入必需的软件包：

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

5. 编译 LiteRT 模型

本部分介绍了一些高级用法，例如直接编译 LiteRT (TFLite) 模型。

从 TFLite 模型进行 EdgeTPU 编译

此步骤需要 TFLite 模型。如果您没有 TFLite 模型，请将模型转换为 TFLite 格式。

获取 TFLite 模型

在此用例中，我们使用 MediaPipe MultiClass Segmentation 模型。

您可以从 MediaPipe 图像分割 页面获取 TFLite 模型。

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)

使用 LiteRT Python API 快速验证 TfLite 模型

在以下示例中，您将看到遮盖图片和混合结果。

# 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 编译。

我们使用 ai_edge_litert.aot 中的 API 来编译模型。

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 方法将所有模型导出到磁盘。

默认情况下，模型以扁平结构存储在输出目录中，每个模型名称都带有后端 ID 后缀。

例如：

模型文件名

后端

SoC

备注

selfie_segmentation_fallback.tflite

CPU/GPU

不适用

不适用

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. 在 CPU 上导出和验证

编译完成后，在 CPU 上验证 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（提前）模块提供了专门用于此目的的 ai_pack 实用程序。这些实用程序会创建一个 AI 包，这是一个至关重要的数据资产。AI 包会将已编译的模型与设备定位配置捆绑在一起，确保将正确的模型和资产交付给适当的用户设备。这对于 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
)

检查 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. 配置高级选项

针对特定设备或 EdgeTPU 进行 NPU 编译

默认情况下，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 图像分割示例。