Triển khai LiteRT trên Google Tensor

1. Tổng quan

Google Tensor SDK được dùng để biên dịch các mô hình LiteRT cho thiết bị Pixel. Bạn có thể triển khai các mô hình đã biên dịch trên thiết bị Pixel để nâng cao hiệu suất suy luận ML. Để sử dụng SDK, trước tiên bạn phải chuyển đổi mô hình của mình thành mô hình LiteRT (tflite).

Lớp học lập trình này dựa trên colab chung trên GitHub, LiteRT AOT Compilation Tutorial Colab:.

Mục tiêu

Tìm hiểu cách sử dụng trình biên dịch LiteRT AOT (trước thời gian chạy) để biên dịch một mô hình phân đoạn ảnh chân dung tự chụp từ mô hình TFLite thành mô hình LiteRT được tối ưu hoá và biên dịch cho EdgeTPU trên thiết bị.

Colab này cũng hướng dẫn bạn các bước chuẩn bị mô hình bằng Play cho AI trên thiết bị (PODAI).

PODAI cung cấp các mô hình tuỳ chỉnh cho các tính năng AI trên thiết bị một cách hiệu quả hơn. Công cụ này giúp đơn giản hoá quy trình ra mắt, nhắm mục tiêu, tạo phiên bản và tải mô hình AI xuống. Khi kết hợp với quá trình biên dịch AOT EdgeTPU LiteRT, phương pháp này cho phép nhà phát triển phân phối các mô hình ML đã biên dịch cho nhiều thiết bị mà không cần biết điện thoại của người dùng cuối có EdgeTPU nào.

Mô hình đã dùng

Các mô hình mà chúng tôi đang sử dụng ban đầu được xuất bản trên Hướng dẫn phân đoạn hình ảnh của MediaPipe. Dưới đây là một số thông tin chi tiết về mô hình được dùng trong lớp học lập trình này:

• SelfieMulticlass: Một mô hình LiteRT lấy hình ảnh của một người, xác định vị trí của các khu vực như tóc, da và quần áo, đồng thời xuất bản đồ phân đoạn hình ảnh cho các mục này.

2. Bắt đầu

Hãy làm theo các bước sau để truy cập và bắt đầu sử dụng Google Tensor SDK:

1. Đăng ký để có quyền truy cập vào Google Tensor SDK. Trước khi tiếp tục, bạn cần chờ email của Google có chứa đường liên kết tải trình bổ trợ trình biên dịch xuống.
2. Tải trình bổ trợ trình biên dịch (litert_plugin_compiler.tar.gz) xuống rồi đặt trình bổ trợ đó vào một thư mục mà bạn chọn.
3. Đặt biến môi trường thành đường dẫn hệ thống cục bộ của tệp đã tải xuống, GOOGLE_TENSOR_SDK_BETA.
   Bạn có thể chạy lệnh này trên thiết bị đầu cuối bash:

   export GOOGLE_TENSOR_SDK_BETA=/path/to/downloaded/compiler
   

   Hoặc bạn có thể chạy lệnh này trong sổ tay Colab:

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

4. Sau đó, hãy chạy lệnh này để cài đặt gói:

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

3. Cài đặt các gói bắt buộc

Bắt đầu bằng cách cài đặt các gói bắt buộc, bao gồm cả ai-edge-litert-nightly (chứa trình biên dịch AOT EdgeTPU) và các thư viện khác mà bạn dùng để chuyển đổi mô hình.

Dùng gói này để cài đặt phần phụ trợ LiteRT cho Google Tensor: ai-edge-litert-sdk-google-tensor.

Sau khi cài đặt các gói, hãy khởi động lại phiên và tiếp tục từ các bước cài đặt. Không lặp lại quá trình cài đặt.

Lưu ý quan trọng: Nếu bạn dự định sử dụng một sổ tay Colab hoặc sổ tay Jupyter, hãy thêm dấu chấm than ! trước mỗi lệnh pip install và pip uninstall. Hãy làm như vậy khi bạn sao chép và dán lệnh pip vào một ô colab.

Nếu dự định thiết lập trên hệ thống của mình, bạn nên sử dụng môi trường ảo Python (venv) và chạy các lệnh này trong môi trường ảo.

Gỡ cài đặt một số gói

Trước đó, hãy gỡ cài đặt tensorflow đi kèm với thời gian chạy colab theo mặc định.

Lưu ý: Trong bước này, bạn có thể gặp phải cảnh báo về việc chưa cài đặt các gói. Bạn có thể yên tâm bỏ qua những thông báo này.

pip uninstall -y tensorflow ai-edge-litert

Cài đặt tất cả các thư viện

Cài đặt phần phụ trợ LiteRT cho Google Tensor

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

Cài đặt các gói còn lại

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

Lưu ý quan trọng: Nếu bạn sử dụng Colab hoặc sổ tay Jupyter:

• Khởi động lại Colab-runtime sau khi quá trình cài đặt hoàn tất.
• Khởi động lại thời gian chạy của Colab bằng cách nhấp vào Thời gian chạy > Khởi động lại phiên trên thanh trình đơn để các gói đã cài đặt được phản ánh chính xác. Bạn chỉ nên khởi động lại sau khi hoàn tất tất cả các quy trình cài đặt.
• Sau khi khởi động lại thời gian chạy, chỉ thực thi các lệnh sau . Bạn không cần chạy lại các bước trước đó.

4. Nhập tất cả thư viện

Tiến hành thực thi chính sau khi quá trình cài đặt hoàn tất.

Đề xuất: Đối với phần này và các phần sau, hãy sử dụng sổ tay Jupyter hoặc thêm các lệnh vào một tập lệnh Python.

Nhập các gói bắt buộc:

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

Lưu ý: Nếu đang sử dụng gói phát hành ổn định ai-edge-litert thay vì ai-edge-litert-nightly, thì bạn sẽ phải đặt biến môi trường này trước tất cả các câu lệnh nhập (cho mỗi lần chạy tập lệnh): GOOGLE_TENSOR_BACKEND_ENABLED=1. Đoạn mã Python cho cùng một mục đích:

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

OR (vỏ)

export GOOGLE_TENSOR_BACKEND_ENABLED=1

5. Biên dịch mô hình LiteRT

Phần này đề cập đến các cách sử dụng nâng cao như biên dịch trực tiếp một mô hình LiteRT (TFLite).

Biên dịch EdgeTPU từ mô hình TFLite

Bước này yêu cầu một mô hình TFLite. Nếu bạn không có mô hình TFLite, hãy chuyển đổi mô hình của bạn sang định dạng TFLite.

Nhận mô hình TFLite

Chúng tôi sử dụng mô hình Phân đoạn nhiều lớp MediaPipe cho trường hợp sử dụng này.

Bạn có thể tải mô hình TFLite xuống từ trang Phân đoạn hình ảnh bằng 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)

Xác minh nhanh mô hình TfLite bằng LiteRT Python API

Trong ví dụ sau, bạn sẽ thấy cả hình ảnh mặt nạ và kết quả được pha trộn.

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

Chuyển đổi sang mô hình LiteRT, với quá trình biên dịch AOT EdgeTPU.

Chúng tôi sử dụng các API từ ai_edge_litert.aot để biên dịch mô hình.

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

Khi quá trình biên dịch hoàn tất, hãy sử dụng phương thức model.export để xuất tất cả các mô hình ra đĩa.

Theo mặc định, các mô hình được lưu trữ trong một cấu trúc phẳng trong thư mục đầu ra, trong đó mỗi tên mô hình có hậu tố là mã nhận dạng phụ trợ.

Ví dụ:

Tên tệp mô hình	Phụ trợ	SoC	Lưu ý
selfie_segmentation_fallback.tflite	CPU/GPU	Không áp dụng	Không áp dụng
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. Xuất và xác thực trên CPU

Sau khi quá trình biên dịch hoàn tất, hãy xác minh mô hình TFLite trên CPU. Hãy thực hiện việc này bằng cách sử dụng "mô hình dự phòng" được tạo trong quá trình biên dịch.

# 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. Xuất mô hình cho PODAI

Sau khi xác minh các mô hình, bước thiết yếu tiếp theo là chuẩn bị triển khai các mô hình đó. Phần này trình bày chi tiết cách đóng gói các mô hình đã biên dịch để tải lên Google Play, cho phép phân phối đến thiết bị của người dùng thông qua khung AI trên thiết bị của Google Play (PODAI).

Mô-đun AOT (Ahead-of-Time) AiEdgeLiteRT cung cấp các tiện ích ai_pack dành riêng cho mục đích này. Các tiện ích này tạo ra một Gói AI, đây là một tài sản dữ liệu quan trọng. Gói AI kết hợp các mô hình đã biên dịch với cấu hình nhắm đến thiết bị, đảm bảo rằng các mô hình và thành phần phù hợp được phân phối đến thiết bị người dùng thích hợp. Điều này đặc biệt quan trọng đối với các bản biên dịch NPU (Đơn vị xử lý thần kinh), vì nó đảm bảo rằng các mô hình được tối ưu hoá cho một Hệ thống trên một chip (SoC) cụ thể chỉ tiếp cận được những thiết bị được trang bị 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
)

Kiểm tra nguồn Gói 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. Định cấu hình các lựa chọn nâng cao

Biên dịch NPU cho thiết bị cụ thể hoặc EdgeTPU

Theo mặc định, quá trình biên dịch AOT LiteRT sẽ biên dịch cho tất cả các phần phụ trợ đã đăng ký. Đối với quá trình phát triển cục bộ, bạn có thể chỉ muốn biên dịch cho các thiết bị cụ thể, chẳng hạn như điện thoại phát triển. Để đạt được điều này, hãy cung cấp các mục tiêu biên dịch một cách rõ ràng.

Ví dụ sau đây sẽ biên dịch thành 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())

Cờ biên dịch cho Google Tensor

Tuỳ chỉnh quy trình biên dịch thông qua cờ biên dịch. Trong đó, cờ này được dùng: google_tensor_truncation_type="half"

Khi biên dịch một mô hình TFLite

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

9. Các bước tiếp theo

Xin chúc mừng!

Các mô hình của bạn đã sẵn sàng để PODAI sử dụng!

Bây giờ, hãy chuyển sang Android Studio để thực hiện các bước sau; hãy xem Các mẫu phân đoạn hình ảnh LiteRT để biết thông tin chi tiết.