[Python, Tensorflow] Convert .h5 to .onnx 및 Inference(.h5 모델을 .onnx 모델로 변환 및 추론) 하는 방법

오늘은 Tensorflow 2 모델(.h5)을 .onnx 모델로 변환하는 방법에 대해서 알아보도록 하겠습니다.

 

import tensorflow as tf
from tensorflow.keras.models import load_model

model = load_model("model.h5")

먼저, .h5 모델을 로드 합니다. 이어서 .onnx 파일로 변환하기 위해 tf2onnx 패키지를 설치합니다.

 

pip install -U tf2onnx

 

패키지를 설치 후 다음 명령어를 통해 .onnx 모델로 변환합니다.

 

python -m tf2onnx.convert --input model.h5 --output model.onnx --opset 13

 

위의 명령어에서 Opset은 ONNX에서 지원하는 연산의 집합을 의미하며, Opset 버전이 ONNX Runtime과 호환이 되는지 유의해야합니다.

 

.onnx 모델을 사용해서 Inference(추론)을 하기 위해서는 Onnx Runtime 패키지를 설치해주어야 합니다.

 

pip install onnxruntime-gpu

 

위의 패키지 설치 후 다음 코드를 통해 추론을 수행합니다.

 

import onnxruntime as ort
import numpy as np

ort_model = ort.InferenceSession('model.onnx', providers=["CUDAExecutionProvider"])

input_name = ort_model.get_inputs()[0].name
output_name = ort_model.get_outputs()[0].name

test_x = np.zeros(...).astype(np.float32)

output = ort_model.run(None, {input_name : test_x})
print(output)

 

지금까지, Tensorflow 2 모델(.h5)를 .onnx 모델로 변환하고 해당 모델을 사용해서 Inference(추론)하는 방법에 대해서 알아보았습니다.