API Python LiteRT CompiledModel

L'API LiteRT CompiledModel est disponible en Python. Elle offre une interface de haut niveau pour compiler et exécuter des modèles TFLite avec le runtime LiteRT.

Le guide suivant montre l'inférence de processeur de base avec l'API Python CompiledModel.

Installer le package pip

Installez le package pip LiteRT dans votre environnement Python :

pip install ai-edge-litert

Inférence de base

Créer CompiledModel

Créez un modèle compilé à partir d'un fichier .tflite. Le wrapper Python actuel est compilé pour le processeur par défaut.

from ai_edge_litert.compiled_model import CompiledModel

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

Vous pouvez également créer un modèle compilé à partir d'un tampon en mémoire :

from ai_edge_litert.compiled_model import CompiledModel

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

Créer des tampons d'entrée et de sortie

Créez les structures de données (tampons) nécessaires pour contenir les données d'entrée que vous fournirez au modèle pour l'inférence, ainsi que les données de sortie que le modèle produira après l'exécution de l'inférence.

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

La valeur signature_index de 0 sélectionne la première signature du modèle.

Si vous utilisez la mémoire du processeur, remplissez les entrées en écrivant directement des tableaux NumPy dans les tampons d'entrée.

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)

Appeler le modèle

Fournissez les tampons d'entrée et de sortie, puis exécutez le modèle.

model.run_by_index(signature_index, input_buffers, output_buffers)

Récupérer les sorties

Récupérez les sorties en lisant directement la sortie du modèle à partir de la mémoire.

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)

Utiliser TensorBuffer

LiteRT offre une compatibilité intégrée pour l'interopérabilité des tampons d'E/S via l'API TensorBuffer, qui permet d'écrire des tableaux NumPy (write) et de lire des tableaux NumPy (read). Les types de données compatibles sont np.float32, np.int32 et np.int8.

Vous pouvez également créer un tampon soutenu par la mémoire hôte existante :

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)

Pour exécuter par nom de signature, commencez par inspecter les signatures du modèle, puis fournissez des mappages des noms d'entrée/sortie aux instances 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},
)

Pour obtenir une vue plus complète de l'implémentation de l'API TensorBuffer, consultez le code source sur TensorBuffer.

Utiliser l'accélérateur GPU

Si vous disposez d'un GPU, vous pouvez l'utiliser en ajoutant simplement l'option HardwareAccelerator.GPU à l'API de création 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)

Consultez cette page pour savoir quel backend est compatible avec votre plate-forme.