API Python скомпилированной модели LiteRT, API Python скомпилированной модели LiteRT, API Python скомпилированной модели LiteRT

API LiteRT CompiledModel доступен на языке Python и предоставляет высокоуровневый интерфейс для компиляции и запуска моделей TFLite с использованием среды выполнения LiteRT.

В этом руководстве показан базовый алгоритм определения параметров ЦП с помощью API CompiledModel на языке Python.

Установите пакет pip.

Установите пакет LiteRT через pip в вашей среде Python:

pip install ai-edge-litert

Базовый вывод

Создать CompiledModel

Создайте скомпилированную модель из файла .tflite . Текущая оболочка Python по умолчанию компилируется для ЦП.

from ai_edge_litert.compiled_model import CompiledModel

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

Также можно создать скомпилированную модель из буфера, находящегося в оперативной памяти:

from ai_edge_litert.compiled_model import CompiledModel

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

Создание входных и выходных буферов

Создайте необходимые структуры данных (буферы) для хранения входных данных, которые вы будете подавать в модель для вывода результатов, и выходных данных, которые модель выдает после выполнения вывода.

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

Значение signature_index , равное 0 выбирает первую подпись в модели.

Если вы используете память ЦП, заполните входные данные, записывая массивы NumPy непосредственно во входные буферы.

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)

Вызовите модель

Предоставив входные и выходные буферы, запустите модель.

model.run_by_index(signature_index, input_buffers, output_buffers)

Получение результатов

Получение результатов осуществляется путем прямого считывания выходных данных модели из памяти.

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

LiteRT обеспечивает встроенную поддержку взаимодействия буферов ввода-вывода через API TensorBuffer , который поддерживает запись массивов NumPy ( write ) и чтение массивов NumPy ( read ). Поддерживаемые типы данных: np.float32 , np.int32 и np.int8 .

Также можно создать буфер, используя существующую память хоста:

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)

Для запуска по имени сигнатуры сначала проверьте сигнатуры модели, а затем передайте карты имен входных/выходных данных экземплярам 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},
)

Более полное представление о реализации API TensorBuffer можно найти в исходном коде по адресу TensorBuffer .

Используйте графический ускоритель (GPU Accelerator)

Если у вас есть графический процессор (GPU), вы можете использовать его, просто добавив опцию HardwareAccelerator.GPU в API создания 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)

Проверьте это , чтобы узнать, какой бэкэнд поддерживается вашей платформой.