LiteRT CompiledModel Python API

LiteRT CompiledModel API 提供 Python 版本，可透過高階介面，使用 LiteRT 執行階段編譯及執行 TFLite 模型。

以下指南說明如何使用 CompiledModel Python API 進行基本的 CPU 推論。

安裝 pip 套件

在 Python 環境中安裝 LiteRT pip 套件：

pip install ai-edge-litert

基本推論

建立「CompiledModel」

從 .tflite 檔案建立已編譯的模型。目前的 Python 包裝函式預設會針對 CPU 編譯。

from ai_edge_litert.compiled_model import CompiledModel

model = CompiledModel.from_file("mymodel.tflite")

您也可以從記憶體內緩衝區建立已編譯模型：

from ai_edge_litert.compiled_model import CompiledModel

with open("mymodel.tflite", "rb") as f:
  model = CompiledModel.from_buffer(f.read())

建立輸入和輸出緩衝區

建立必要的資料結構 (緩衝區)，以保存您將提供給模型進行推論的輸入資料，以及模型在執行推論後產生的輸出資料。

signature_index = 0
input_buffers = model.create_input_buffers(signature_index)
output_buffers = model.create_output_buffers(signature_index)

signature_index 的 0 值會選取模型中的第一個簽章。

如果您使用 CPU 記憶體，請直接將 NumPy 陣列寫入輸入緩衝區，填寫輸入內容。

import numpy as np

input_data = np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32)
input_buffers[0].write(input_data)

叫用模型

提供輸入和輸出緩衝區，然後執行模型。

model.run_by_index(signature_index, input_buffers, output_buffers)

擷取輸出內容

直接從記憶體讀取模型輸出，即可擷取輸出內容。

import numpy as np

# Replace num_elements with the size of your model's output tensor.
num_elements = 4
output_array = output_buffers[0].read(num_elements, np.float32)

使用TensorBuffer

LiteRT 透過 TensorBuffer API 內建支援 I/O 緩衝區互通性，可支援寫入 NumPy 陣列 (write) 和讀取 NumPy 陣列 (read)。支援的 Dtype 為 np.float32、np.int32 和 np.int8。

您也可以建立由現有主機記憶體支援的緩衝區：

import numpy as np
from ai_edge_litert.tensor_buffer import TensorBuffer

input_array = np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32)
input_buffer = TensorBuffer.create_from_host_memory(input_array)

如要依簽章名稱執行，請先檢查模型簽章，然後從輸入/輸出名稱提供對應至 TensorBuffer 執行個體的對映：

from ai_edge_litert.tensor_buffer import TensorBuffer

signatures = model.get_signature_list()
# Example signature structure:
# {"serving_default": {"inputs": ["input_0"], "outputs": ["output_0"]}}

input_buffer = TensorBuffer.create_from_host_memory(input_array)
output_buffer = model.create_output_buffer_by_name("serving_default", "output_0")

model.run_by_name(
  "serving_default",
  {"input_0": input_buffer},
  {"output_0": output_buffer},
)

如要更全面瞭解 TensorBuffer API 的實作方式，請參閱 TensorBuffer 的原始碼。

使用 GPU 加速器

如果您有 GPU，只要將 HardwareAccelerator.GPU 選項新增至 CompiledModel 建立 API，即可使用 GPU。

from ai_edge_litert.compiled_model import CompiledModel
from ai_edge_litert.compiled_model import HardwareAccelerator

model = CompiledModel.from_file("mymodel.tflite", HardwareAccelerator.GPU)

請參閱這篇文章，瞭解平台支援哪些後端。