ONNX Runtime 파이썬 API — InferenceSession과 IOBinding

ONNX Runtime 파이썬 API — InferenceSession과 IOBinding

ONNX Runtime 파이썬 API의 중심에는 InferenceSession 클래스가 있고, 데이터 입출력 관점에서는 OrtValue와 IOBinding이 핵심 개념이에요.

출처: https://onnxruntime.ai/docs/api/python/api_summary.html

InferenceSession

모델 실행의 메인 클래스예요. 기본은 ONNX 포맷이고, .ort 확장자는 ORT 포맷으로 인식됩니다.

import onnxruntime
session = onnxruntime.InferenceSession('model.onnx',
    providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])

주요 메서드:

  • get_inputs(), get_outputs() — 입·출력 메타데이터를 NodeArg 목록으로 반환
  • get_modelmeta() — 모델 메타데이터(ModelMetadata) 반환
  • run(output_names, input_feed) — 추론 실행
  • disable_fallback() / enable_fallback() — 세션 run 실패 시 폴백 제어

OrtValue와 IOBinding

기본적으로 입·출력은 CPU에 놓입니다. GPU처럼 디바이스 위에서 데이터를 직접 다루려면 IOBinding을 써요.

session = onnxruntime.InferenceSession('model.onnx',
    providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])

io_binding = session.io_binding()
io_binding.bind_cpu_input('input', X)   # CPU 입력 복사
io_binding.bind_output('output')        # 출력 바인딩
session.run_with_iobinding(io_binding)
Y = io_binding.copy_outputs_to_cpu()[0]

run_with_iobinding()은 run() 대신 IOBinding을 쓰는 방식이에요. 데이터를 GPU에 직접 올려 CPU↔GPU 복사를 줄이면 대용량 입출력에서 큰 이득을 볼 수 있습니다. OrtValue.ortvalue_from_numpy(X, 'cuda', 0)처럼 OrtValue를 만들어 bind_input에 넘기는 방식도 있고, 동적 shape 출력은 bind_output('output', 'cuda') 형태로 ONNX Runtime이 디바이스 메모리를 직접 할당하게 할 수 있어요.

더 알아보기