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)

0signature_index 值会选择模型中的第一个签名。

如果您使用的是 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 缓冲区互操作性提供内置支持,该 API 支持写入 NumPy 数组 (write) 和读取 NumPy 数组 (read)。支持的 dtype 为 np.float32np.int32np.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,只需向 CompiledModel 创建 API 添加 HardwareAccelerator.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)

请点击此处查看您的平台支持哪个后端。