API de Python de LiteRT CompiledModel

La API de LiteRT CompiledModel está disponible en Python y ofrece una interfaz de alto nivel para compilar y ejecutar modelos de TFLite con el tiempo de ejecución de LiteRT.

En la siguiente guía, se muestra la inferencia básica de la CPU con la API de CompiledModel de Python.

Instala el paquete pip

Instala el paquete pip de LiteRT en tu entorno de Python:

pip install ai-edge-litert

Inferencia básica

Crear CompiledModel

Crea un modelo compilado a partir de un archivo .tflite. El wrapper de Python actual se compila para CPU de forma predeterminada.

from ai_edge_litert.compiled_model import CompiledModel

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

También puedes crear un modelo compilado a partir de un búfer en la memoria:

from ai_edge_litert.compiled_model import CompiledModel

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

Cómo crear búferes de entrada y salida

Crea las estructuras de datos (búferes) necesarias para contener los datos de entrada que ingresarás en el modelo para la inferencia y los datos de salida que el modelo produce después de ejecutar la inferencia.

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

El valor signature_index de 0 selecciona la primera firma del modelo.

Si usas la memoria de la CPU, completa las entradas escribiendo arrays de NumPy directamente en los búferes de entrada.

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)

Invoca el modelo

Proporciona los búferes de entrada y salida, y ejecuta el modelo.

model.run_by_index(signature_index, input_buffers, output_buffers)

Recupera resultados

Recupera los resultados leyendo directamente el resultado del modelo de la memoria.

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)

Usa TensorBuffer

LiteRT proporciona compatibilidad integrada para la interoperabilidad de búferes de E/S a través de la API de TensorBuffer, que admite la escritura de arrays de NumPy (write) y la lectura de arrays de NumPy (read). Los dtypes admitidos son np.float32, np.int32 y np.int8.

También puedes crear un búfer respaldado por la memoria del host existente:

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)

Para ejecutar por nombre de firma, primero inspecciona las firmas del modelo y, luego, proporciona mapas de los nombres de entrada o salida a las instancias de 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},
)

Para obtener una vista más completa de cómo se implementa la API de TensorBuffer, consulta el código fuente en TensorBuffer.

Usar acelerador de GPU

Si tienes una GPU, puedes usarla con solo agregar la opción HardwareAccelerator.GPU a la API de creación de CompiledModel.

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

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

Consulta este vínculo para ver qué backend es compatible con tu plataforma.