keras h5 模型轉換 onnx 并推理測試

目錄

• • 一.版本介紹
  • 二.轉換流程
  • 三.轉換過程
  • 四.推理測試

一.版本介紹

轉換：

keras 2.1.5

tensorflow 1.13.1

tf2onnx 1.5.5

推理:

opencv 4.4.0

onnx 1.5.0

onnxruntime 1.6.0

二.轉換流程

• ① h5 to pb
• ② pb to onnx

三.轉換過程

首先準備自己的h5模型；這裡要注意h5模型的儲存方式，不同儲存方式，對應不同加載方式，如下：

> 訓練:  model.save()    
> 加載:  model.load_model()
           

> 訓練:  model.save_weights()    
> 加載:    
> 		model = LeNet()     
> 		model.load_weights(weight_path) 
           

• ① h5 to pb

# h5 to  pb
from keras.models import load_model
import tensorflow as tf
import os
from keras import backend as K
from tensorflow.python.framework import graph_util, graph_io
import tensorflow as tf
from tensorflow.python.platform import gfile

def h5_to_pb(h5_weight_path, output_dir, out_prefix="output_", log_tensorboard=True):
    if not os.path.exists(output_dir):
        os.mkdir(output_dir)
    h5_model =  load_model(h5_weight_path) 
    out_nodes = []
    for i in range(len(h5_model.outputs)):
        out_nodes.append(out_prefix + str(i + 1))
        tf.identity(h5_model.output[i], out_prefix + str(i + 1))

    model_name = os.path.splitext(os.path.split(h5_weight_path)[-1])[0] + '.pb'

    sess = K.get_session()
    init_graph = sess.graph.as_graph_def()
    main_graph = graph_util.convert_variables_to_constants(sess, init_graph, out_nodes)
    graph_io.write_graph(main_graph, output_dir, name=model_name, as_text=False)

#讀取模型各層
def read_pb(GRAPH_PB_PATH):
    with tf.Session() as sess:
        print("load graph！！！")
        with gfile.FastGFile(GRAPH_PB_PATH, 'rb') as f:
            graph_def = tf.GraphDef()在這裡插入代碼片
            graph_def.ParseFromString(f.read())
            tf.import_graph_def(graph_def, name='')
            for i, n in enumerate(graph_def.node):
                print("Name of the node - %s" % n.name)

h5_to_pb(h5_weight_path='model.h5', output_dir='./')
read_pb('./model.pb')
           

• ② pb to onnx

#安裝tf2onnx
pip install tf2onnx==1.5.5
           

這裡需要使用netron檢視pb模型的輸入，輸出層名稱。

https://netron.app/

python -m tf2onnx.convert --input ./model.pb  --inputs conv2d_1_input:0 --outputs output_1:0 --output ./model.onnx
           

四.推理測試

這裡使用minst資料集訓練的lenet模型測試，故圖像預處理為(28,28,1) 。

import cv2
from keras.models import load_model
import numpy as np
import onnxruntime as rt
from keras.preprocessing.image import img_to_array
import time

#keras推理
img = cv2.imread('1.jpg',0)
img = cv2.resize(img,(28,28))
img = np.array([img_to_array(img)],dtype='float')/255.0
model = load_model('model.h5')
predict = model.predict(img)
p = [round(i,5) for i in predict[0]]
print('\n keras :')
print(p)

#opencv dnn pb
img = cv2.imread('1.jpg',0)
img = cv2.resize(img,(28,28))
img = np.array([img_to_array(img)],dtype='float')/255.0
img = img.transpose((0,3,1,2))
net = cv2.dnn.readNetFromTensorflow("model.pb")
net.setInput(img)  
out = net.forward()  
print('\n opencv dnn pb:')
p = [round(i,5) for i in out[0]]
print(p)

#onnxruntime 
img = cv2.imread('1.jpg',0)
img = cv2.resize(img,(28,28))
img = np.array([img_to_array(img)],dtype='float')/255.0
img= img.astype(np.float32)
sess=rt.InferenceSession('model.onnx')
input_name=sess.get_inputs()[0].name
res=sess.run(None,{input_name:img})
print('\n onnxruntime :')
print([round(i,5) for i in res[0]])
           

如圖推理結果一緻即可: