API LiteRT CompiledModel có trong Python, cung cấp một giao diện cấp cao để biên dịch và chạy các mô hình TFLite bằng thời gian chạy LiteRT.
Hướng dẫn sau đây trình bày suy luận cơ bản về CPU bằng CompiledModel Python API.
Cài đặt gói pip
Cài đặt gói pip LiteRT trong môi trường Python:
pip install ai-edge-litert
Suy luận cơ bản
Tạo CompiledModel
Tạo một mô hình đã biên dịch từ tệp .tflite. Trình bao bọc Python hiện tại sẽ biên dịch cho CPU theo mặc định.
from ai_edge_litert.compiled_model import CompiledModel
model = CompiledModel.from_file("mymodel.tflite")
Bạn cũng có thể tạo một mô hình đã biên dịch từ vùng đệm trong bộ nhớ:
from ai_edge_litert.compiled_model import CompiledModel
with open("mymodel.tflite", "rb") as f:
model = CompiledModel.from_buffer(f.read())
Tạo vùng đệm đầu vào và đầu ra
Tạo các cấu trúc dữ liệu (vùng đệm) cần thiết để lưu trữ dữ liệu đầu vào mà bạn sẽ đưa vào mô hình để suy luận, cũng như dữ liệu đầu ra mà mô hình tạo ra sau khi chạy suy luận.
signature_index = 0
input_buffers = model.create_input_buffers(signature_index)
output_buffers = model.create_output_buffers(signature_index)
Giá trị signature_index của 0 sẽ chọn chữ ký đầu tiên trong mô hình.
Nếu bạn đang sử dụng bộ nhớ CPU, hãy điền các đầu vào bằng cách ghi trực tiếp các mảng NumPy vào các vùng đệm đầu vào.
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)
Gọi mô hình
Cung cấp vùng đệm đầu vào và đầu ra, chạy mô hình.
model.run_by_index(signature_index, input_buffers, output_buffers)
Truy xuất dữ liệu đầu ra
Truy xuất đầu ra bằng cách đọc trực tiếp đầu ra của mô hình từ bộ nhớ.
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)
Sử dụng TensorBuffer
LiteRT cung cấp khả năng hỗ trợ tích hợp cho khả năng tương tác của vùng đệm I/O thông qua API TensorBuffer, hỗ trợ ghi mảng NumPy (write) và đọc mảng NumPy (read). Các dtypes được hỗ trợ là np.float32, np.int32 và np.int8.
Bạn cũng có thể tạo một vùng đệm được hỗ trợ bởi bộ nhớ máy chủ hiện có:
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)
Để chạy theo tên chữ ký, trước tiên, hãy kiểm tra các chữ ký mô hình, sau đó cung cấp các bản đồ từ tên đầu vào/đầu ra đến các thực thể 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},
)
Để biết thông tin đầy đủ hơn về cách triển khai TensorBuffer API, hãy xem mã nguồn tại TensorBuffer.
Sử dụng Trình tăng tốc GPU
Nếu có GPU, bạn có thể sử dụng GPU chỉ bằng cách thêm lựa chọn HardwareAccelerator.GPU vào API tạo 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)
Hãy kiểm tra nội dung này để biết nền tảng của bạn hỗ trợ những phần phụ trợ nào.