LiteRT CompiledModel Python API

Die LiteRT CompiledModel API ist in Python verfügbar und bietet eine übergeordnete Schnittstelle zum Kompilieren und Ausführen von TFLite-Modellen mit der LiteRT-Laufzeit.

Im folgenden Leitfaden wird die grundlegende CPU-Inferenz mit der CompiledModel-Python-API beschrieben.

pip-Paket installieren

Installieren Sie das LiteRT-PIP-Paket in Ihrer Python-Umgebung:

pip install ai-edge-litert

Einfache Inferenz

CompiledModel erstellen

Erstellen Sie ein kompiliertes Modell aus einer .tflite-Datei. Der aktuelle Python-Wrapper wird standardmäßig für die CPU kompiliert.

from ai_edge_litert.compiled_model import CompiledModel

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

Sie können auch ein kompiliertes Modell aus einem In-Memory-Puffer erstellen:

from ai_edge_litert.compiled_model import CompiledModel

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

Eingabe- und Ausgabepuffer erstellen

Erstellen Sie die erforderlichen Datenstrukturen (Puffer) für die Eingabedaten, die Sie in das Modell für die Inferenz einfügen, und die Ausgabedaten, die das Modell nach der Ausführung der Inferenz erzeugt.

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

Der signature_index-Wert von 0 wählt die erste Signatur im Modell aus.

Wenn Sie CPU-Arbeitsspeicher verwenden, füllen Sie die Eingaben, indem Sie NumPy-Arrays direkt in die Eingabepuffer schreiben.

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)

Modell aufrufen

Stellen Sie die Ein- und Ausgabepuffer bereit und führen Sie das Modell aus.

model.run_by_index(signature_index, input_buffers, output_buffers)

Ausgaben abrufen

Ausgaben abrufen, indem die Modellausgabe direkt aus dem Speicher gelesen wird.

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“ verwenden

LiteRT bietet integrierte Unterstützung für die Interoperabilität von I/O-Puffern über die TensorBuffer API, die das Schreiben von NumPy-Arrays (write) und das Lesen von NumPy-Arrays (read) unterstützt. Die unterstützten DTypes sind np.float32, np.int32 und np.int8.

Sie können auch einen Puffer erstellen, der durch vorhandenen Hostspeicher gesichert wird:

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)

Wenn Sie die Ausführung nach Signaturnamen durchführen möchten, prüfen Sie zuerst die Modellsignaturen und stellen Sie dann Zuordnungen von Ein-/Ausgabenamen zu TensorBuffer-Instanzen bereit:

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},
)

Eine umfassendere Ansicht der Implementierung der TensorBuffer API finden Sie im Quellcode unter TensorBuffer.

GPU-Beschleuniger verwenden

Wenn Sie eine GPU haben, können Sie sie verwenden, indem Sie der API zum Erstellen von CompiledModel die Option HardwareAccelerator.GPU hinzufügen.

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

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

Hier finden Sie Informationen dazu, welches Backend für Ihre Plattform unterstützt wird.