LiteRT CompiledModel Python API

API CompiledModel LiteRT tersedia di Python, yang menawarkan antarmuka tingkat tinggi untuk mengompilasi dan menjalankan model TFLite dengan runtime LiteRT.

Panduan berikut menunjukkan inferensi CPU dasar dengan CompiledModel Python API.

Instal paket pip

Instal paket pip LiteRT di lingkungan Python Anda:

pip install ai-edge-litert

Inferensi dasar

Buat CompiledModel

Buat model yang dikompilasi dari file .tflite. Wrapper Python saat ini dikompilasi untuk CPU secara default.

from ai_edge_litert.compiled_model import CompiledModel

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

Anda juga dapat membuat model yang dikompilasi dari buffer dalam memori:

from ai_edge_litert.compiled_model import CompiledModel

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

Membuat Buffer Input dan Output

Buat struktur data (buffer) yang diperlukan untuk menyimpan data input yang akan dimasukkan ke dalam model untuk inferensi, dan data output yang dihasilkan model setelah menjalankan inferensi.

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

Nilai signature_index dari 0 memilih tanda tangan pertama dalam model.

Jika Anda menggunakan memori CPU, isi input dengan menulis array NumPy langsung ke buffer 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)

Memanggil model

Menyediakan buffer input dan output, jalankan model.

model.run_by_index(signature_index, input_buffers, output_buffers)

Mengambil Output

Mengambil output dengan membaca output model langsung dari memori.

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)

Gunakan TensorBuffer

LiteRT menyediakan dukungan bawaan untuk interoperabilitas buffer I/O melalui TensorBuffer API, yang mendukung penulisan array NumPy (write) dan pembacaan array NumPy (read). Dtype yang didukung adalah np.float32, np.int32, dan np.int8.

Anda juga dapat membuat buffer yang didukung oleh memori host yang ada:

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)

Untuk menjalankan menurut nama tanda tangan, periksa tanda tangan model terlebih dahulu, lalu berikan peta dari nama input/output ke instance 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},
)

Untuk melihat gambaran yang lebih lengkap tentang cara penerapan TensorBuffer API, lihat kode sumber di TensorBuffer.

Menggunakan Akselerator GPU

Jika memiliki GPU, Anda dapat menggunakannya hanya dengan menambahkan opsi HardwareAccelerator.GPU ke API pembuatan 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)

Periksa ini untuk melihat backend mana yang didukung untuk platform Anda.