API Python LiteRT CompiledModel

A API LiteRT CompiledModel está disponível em Python, oferecendo uma interface de alto nível para compilar e executar modelos do TFLite com o ambiente de execução LiteRT.

O guia a seguir mostra a inferência básica de CPU com a API CompiledModel Python.

Instalar o pacote pip

Instale o pacote pip LiteRT no seu ambiente Python:

pip install ai-edge-litert

Inferência básica

Criar CompiledModel

Crie um modelo compilado com base em um arquivo .tflite. O wrapper Python atual é compilado para CPU por padrão.

from ai_edge_litert.compiled_model import CompiledModel

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

Também é possível criar um modelo compilado de um buffer na memória:

from ai_edge_litert.compiled_model import CompiledModel

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

Criar buffers de entrada e saída

Crie as estruturas de dados (buffers) necessárias para armazenar os dados de entrada que você vai inserir no modelo para inferência e os dados de saída que o modelo produz após a execução da inferência.

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

O valor signature_index de 0 seleciona a primeira assinatura no modelo.

Se você estiver usando a memória da CPU, preencha as entradas gravando matrizes NumPy diretamente nos buffers 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)

Invocar o modelo

Forneça os buffers de entrada e saída e execute o modelo.

model.run_by_index(signature_index, input_buffers, output_buffers)

Recuperar saídas

Recuperar saídas lendo diretamente a saída do modelo da memória.

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)

Usar TensorBuffer

O LiteRT oferece suporte integrado à interoperabilidade de buffer de E/S pela API TensorBuffer, que permite gravar (write) e ler (read) arrays do NumPy. Os dtypes compatíveis são np.float32, np.int32 e np.int8.

Também é possível criar um buffer com o apoio da memória do host:

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 executar por nome de assinatura, primeiro inspecione as assinaturas do modelo e forneça mapas de nomes de entrada/saída para instâncias 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 uma visão mais completa de como a API TensorBuffer é implementada, consulte o código-fonte em TensorBuffer.

Usar acelerador de GPU

Se você tiver uma GPU, basta adicionar a opção HardwareAccelerator.GPU à API de criação do 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)

Confira este link para saber qual back-end é compatível com sua plataforma.