ONNX Runtime Python 시작하기

ONNX Runtime Python 시작하기

ONNX Runtime을 Python에서 쓰는 첫 관문이에요. CPU는 onnxruntime, GPU(CUDA 12.x)는 onnxruntime-gpu 패키지를 설치하면 돼요. 두 패키지는 한 환경에 동시에 설치하지 않는 게 원칙이에요.

설치

# CPU
pip install onnxruntime

# GPU (CUDA 12.x)
pip install onnxruntime-gpu

모델을 ONNX로 내보낼 때는 파이토치는 torch에 내장된 ONNX 내보내기가, 텐서플로는 tf2onnx, scikit-learn은 skl2onnx가 쓰여요.

추론 기본 흐름

import onnxruntime as ort

sess = ort.InferenceSession("fashion_mnist_model.onnx")
outputs = sess.run(None, {"input": input_numpy})

확인 포인트

  • sess.run(None, ...)의 첫 인자는 가져올 출력 이름 목록이고, None이면 모든 출력을 돌려줘요.
  • 입력은 numpy 배열 같은 배열형이어야 해요.

출처: https://onnxruntime.ai/docs/get-started/with-python.html

더 알아보기