API Python LiteRT CompiledModel

L'API LiteRT CompiledModel è disponibile in Python e offre un'interfaccia di alto livello per compilare ed eseguire modelli TFLite con il runtime LiteRT.

La seguente guida mostra l'inferenza di base della CPU con l'API Python CompiledModel.

Installa il pacchetto pip

Installa il pacchetto pip LiteRT nel tuo ambiente Python:

pip install ai-edge-litert

Inferenza di base

Crea CompiledModel

Crea un modello compilato da un file .tflite. Il wrapper Python attuale viene compilato per la CPU per impostazione predefinita.

from ai_edge_litert.compiled_model import CompiledModel

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

Puoi anche creare un modello compilato da un buffer in memoria:

from ai_edge_litert.compiled_model import CompiledModel

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

Creare buffer di input e output

Crea le strutture di dati (buffer) necessarie per contenere i dati di input che inserirai nel modello per l'inferenza e i dati di output che il modello produce dopo l'esecuzione dell'inferenza.

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

Il valore signature_index di 0 seleziona la prima firma nel modello.

Se utilizzi la memoria della CPU, compila gli input scrivendo direttamente gli array NumPy nei buffer di input.

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)

Richiamare il modello

Fornisci i buffer di input e output ed esegui il modello.

model.run_by_index(signature_index, input_buffers, output_buffers)

Recuperare gli output

Recupera gli output leggendo direttamente l'output del modello dalla 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)

Utilizza TensorBuffer

LiteRT fornisce il supporto integrato per l'interoperabilità dei buffer I/O tramite l'API TensorBuffer, che supporta la scrittura di array NumPy (write) e la lettura di array NumPy (read). I tipi di dati supportati sono np.float32, np.int32 e np.int8.

Puoi anche creare un buffer supportato dalla memoria host esistente:

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)

Per l'esecuzione in base al nome della firma, ispeziona prima le firme del modello e poi fornisci mappe dai nomi di input/output alle istanze 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},
)

Per una visione più completa dell'implementazione dell'API TensorBuffer, consulta il codice sorgente all'indirizzo TensorBuffer.

Utilizza l'acceleratore GPU

Se hai una GPU, puoi utilizzarla semplicemente aggiungendo l'opzione HardwareAccelerator.GPU all'API di creazione di 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 questa pagina per scoprire quale backend è supportato per la tua piattaforma.