LiteRT CompiledModel Python API

LiteRT CompiledModel API는 Python에서 사용할 수 있으며, LiteRT 런타임으로 TFLite 모델을 컴파일하고 실행하기 위한 고급 인터페이스를 제공합니다.

다음 가이드에서는 CompiledModel Python API를 사용한 기본 CPU 추론을 보여줍니다.

pip 패키지 설치

Python 환경에 LiteRT pip 패키지를 설치합니다.

pip install ai-edge-litert

기본 추론

CompiledModel 만들기

.tflite 파일에서 컴파일된 모델을 만듭니다. 현재 Python 래퍼는 기본적으로 CPU용으로 컴파일됩니다.

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)

0의 signature_index 값은 모델의 첫 번째 서명을 선택합니다.

CPU 메모리를 사용하는 경우 입력 버퍼에 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는 NumPy 배열 쓰기 (write) 및 NumPy 배열 읽기 (read)를 지원하는 TensorBuffer API를 통해 I/O 버퍼 상호 운용성을 기본적으로 지원합니다. 지원되는 dtypes는 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},
)

TensorBuffer API가 구현되는 방식을 더 자세히 알아보려면 TensorBuffer의 소스 코드를 참고하세요.

GPU 가속기 사용

GPU가 있는 경우 컴파일된 모델 생성 API에 HardwareAccelerator.GPU 옵션을 추가하여 사용할 수 있습니다.

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

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

여기에서 플랫폼에 지원되는 백엔드를 확인하세요.