Keras 모델( .h5)을 onnx로 변환 및 OpenCV DNN에서 사용하기

Keras 모델( .h5)을 onnx로 변환한 후,  OpenCV DNN에서 onnx 모델을 사용해서 추론을 해봅니다.

 

2022. 10. 25  최초작성

1. MNIST 데이터셋에 대한 모델을 학습시켜 h5 모델 파일로 저장합니다. 

코드 출처 https://www.tensorflow.org/tutorials/quickstart/beginner?hl=ko 

 

import tensorflow as tf mnist = tf.keras.datasets.mnist (x_train, y_train), (x_test, y_test) = mnist.load_data() x_train, x_test = x_train / 255.0, x_test / 255.0 model = tf.keras.models.Sequential([   tf.keras.layers.Flatten(input_shape=(28, 28)),   tf.keras.layers.Dense(128, activation='relu'),   tf.keras.layers.Dropout(0.2),   tf.keras.layers.Dense(10, activation='softmax') ]) model.compile(optimizer='adam',               loss='sparse_categorical_crossentropy',               metrics=['accuracy']) model.fit(x_train, y_train, epochs=10) model.evaluate(x_test,  y_test, verbose=2) # 모델을 저장합니다. model.save('mnist.h5')

Epoch 1/10 1875/1875 [==============================] - 15s 5ms/step - loss: 0.2943 - accuracy: 0.9141 Epoch 2/10 1875/1875 [==============================] - 9s 5ms/step - loss: 0.1393 - accuracy: 0.9585 Epoch 3/10 1875/1875 [==============================] - 9s 5ms/step - loss: 0.1059 - accuracy: 0.9674 Epoch 4/10 1875/1875 [==============================] - 9s 5ms/step - loss: 0.0883 - accuracy: 0.9729 Epoch 5/10 1875/1875 [==============================] - 16s 8ms/step - loss: 0.0745 - accuracy: 0.9767 Epoch 6/10 1875/1875 [==============================] - 16s 9ms/step - loss: 0.0670 - accuracy: 0.9786 Epoch 7/10 1875/1875 [==============================] - 13s 7ms/step - loss: 0.0597 - accuracy: 0.9807 Epoch 8/10 1875/1875 [==============================] - 13s 7ms/step - loss: 0.0522 - accuracy: 0.9833 Epoch 9/10 1875/1875 [==============================] - 15s 8ms/step - loss: 0.0486 - accuracy: 0.9838 Epoch 10/10 1875/1875 [==============================] - 15s 8ms/step - loss: 0.0429 - accuracy: 0.9859 313/313 - 4s - loss: 0.0797 - accuracy: 0.9779 - 4s/epoch - 11ms/step

 

2. h5 모델 파일을 onnx 모델 파일로 변환합니다. 

참고 https://onnxruntime.ai/docs/tutorials/tf-get-started.html 

 

우선 tf2onnx 패키지를 설치해야 합니다.

pip install tf2onnx

 

다음 코드에서 h5 모델 파일을 pb 모델 파일로 변환 한후, 다시 onnx 모델 파일로 변환합니다.

 

# h5 to pb import tensorflow as tf model = tf.keras.models.load_model('mnist.h5', compile=False) model.save('./model', save_format="tf") # pb to onnx import os os.system('python -m tf2onnx.convert --saved-model ./model/ --output model.onnx  --opset 13')

3. OpenCV DNN을 사용하여 onnx 모델로 추론을 해봅니다.

참고 https://jeanvitor.com/how-to-load-pytorch-models-with-opencv/ 

 

import cv2 import numpy as np import tensorflow as tf onnx_model_path = "model.onnx" net =  cv2.dnn.readNetFromONNX(onnx_model_path) mnist = tf.keras.datasets.mnist (x_train, y_train), (x_test, y_test) = mnist.load_data() # 정규화를 해주면 안됩니다. # x_train, x_test = x_train / 255.0, x_test / 255.0 for index in range(5):     blob = cv2.dnn.blobFromImage(x_test[index],  # 추론할 입력 이미지입니다.                                 1.0 / 255,   # 학습할때 이미지의 픽셀값을 255로 나누어 정규화 했습니다.                                 (28, 28),    # 이미지 크기입니다.                                 (0, 0, 0),                                 swapRB=True, # BGR 또는 RGB인지에 따라 설정합니다. 흑백에선 상관없습니다.                                 crop=False)     net.setInput(blob)     preds = net.forward()     biggest_pred_index = np.array(preds)[0].argmax()     print ("예측 ",biggest_pred_index, '  정답 ', y_test[index])

실행 결과 모델을 사용하여 예측한 값과 실제 정답이 출력됩니다.

예측  7   정답  7 예측  2   정답  2 예측  1   정답  1 예측  0   정답  0 예측  4   정답  4