Interfejs API w Pythonie dla skompilowanego modelu LiteRT

Interfejs LiteRT CompiledModel API jest dostępny w języku Python i oferuje interfejs wysokiego poziomu do kompilowania i uruchamiania modeli TFLite w środowisku wykonawczym LiteRT.

W tym przewodniku znajdziesz podstawowe informacje o wnioskowaniu na procesorze za pomocą interfejsu CompiledModel Python API.

Instalowanie pakietu pip

Zainstaluj pakiet pip LiteRT w środowisku Pythona:

pip install ai-edge-litert

Podstawowe wnioskowanie

Utwórz: CompiledModel

Utwórz skompilowany model z pliku .tflite. Obecny moduł Pythona jest domyślnie kompilowany na potrzeby procesora.

from ai_edge_litert.compiled_model import CompiledModel

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

Model skompilowany możesz też utworzyć z bufora w pamięci:

from ai_edge_litert.compiled_model import CompiledModel

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

Tworzenie buforów wejściowych i wyjściowych

Utwórz niezbędne struktury danych (bufory) do przechowywania danych wejściowych, które będziesz przekazywać do modelu na potrzeby wnioskowania, oraz danych wyjściowych, które model generuje po uruchomieniu wnioskowania.

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

Wartość signature_index w przypadku 0 wybiera pierwszy podpis w modelu.

Jeśli używasz pamięci procesora, wypełnij pola, wpisując bezpośrednio tablice NumPy do buforów wejściowych.

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)

Wywoływanie modelu

Podaj bufory wejściowe i wyjściowe, a następnie uruchom model.

model.run_by_index(signature_index, input_buffers, output_buffers)

Pobieranie danych wyjściowych

Pobieraj dane wyjściowe, odczytując je bezpośrednio z pamięci.

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)

Używaj klawisza TensorBuffer

LiteRT zapewnia wbudowaną obsługę współdziałania buforów wejścia/wyjścia za pomocą interfejsu API TensorBuffer, który obsługuje zapisywanie tablic NumPy (write) i odczytywanie tablic NumPy (read). Obsługiwane typy danych to np.float32, np.int32np.int8.

Możesz też utworzyć bufor oparty na istniejącej pamięci hosta:

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)

Aby uruchomić model według nazwy sygnatury, najpierw sprawdź sygnatury modelu, a potem podaj mapy z nazwami wejścia/wyjścia do instancji 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},
)

Aby uzyskać pełniejszy obraz implementacji interfejsu TensorBuffer API, zapoznaj się z kodem źródłowym w TensorBuffer.

Używanie akceleratora GPU

Jeśli masz GPU, możesz go używać, dodając opcję HardwareAccelerator.GPU do interfejsu API tworzenia 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)

Sprawdź tutaj, które zaplecze jest obsługiwane na Twojej platformie.