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)
0 的 signature_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.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,只需向 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)
请点击此处查看您的平台支持哪个后端。