项目中遇到需要将训练好的keras模型转换成onnx以便部署到嵌入式设备进行RT加速，最开始使用的keras2onnx工具，然而此工具支持性并不好，在转化过程中遇到许多问题。因此决定将keras转成tensorflow格式，再使用支持性较好的tf2onnx工具进行转化。

新版本

目前keras2onnx停止更新，功能已整合进tf2onnx中，keras模型可直接使用 api tf2onnx.convert.from_keras()完成转换操作

转换步骤：

1
2
3
4
5
6
7
8
9
10
11
import tf2onnx
import onnxruntime as rt

# 加载模型
model = ResNet50(weights='imagenet')
# 设置输入大小
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)
# 保存onnx路径
output_path = model.name + ".onnx"
# 转换
tf2onnx.convert.from_keras(model, input_signature=spec, opset=13, output_path=output_path)

API说明：

import tf2onnx

model_proto, external_tensor_storage = tf2onnx.convert.from_keras(model,
input_signature=None, opset=None, custom_ops=None,
custom_op_handlers=None, custom_rewriter=None,
inputs_as_nchw=None, extra_opset=None shape_override=None,
target=None, large_model=False, output_path=None)

Args:
    model: the tf.keras model we want to convert
    input_signature: a tf.TensorSpec or a numpy array defining the shape/dtype of the input
    opset: the opset to be used for the ONNX model, default is the latest
    custom_ops: if a model contains ops not recognized by onnx runtime,
        you can tag these ops with a custom op domain so that the
        runtime can still open the model. Type is a dictionary `{op name: domain}`.
    target: list of workarounds applied to help certain platforms
    custom_op_handlers: dictionary of custom ops handlers
    custom_rewriter: list of custom graph rewriters
    extra_opset: list of extra opset's, for example the opset's used by custom ops
    shape_override: dict with inputs that override the shapes given by tensorflow
    inputs_as_nchw: transpose inputs in list from nchw to nhwc
    large_model: use the ONNX external tensor storage format
    output_path: save model to output_path

Returns:
    An ONNX model_proto and an external_tensor_storage dict.

旧版本

一.所需工具

1
2
3
pip install tf2onnx
pip install onnx
pip install onnxruntime

二.转换流程

① h5 to pb
② pb to onnx

三.转换过程

首先准备自己的h5模型；这里要注意h5模型的保存方式，不同保存方式，对应不同加载方式，如下：

1
2
训练:  model.save()    
加载:  model.load_model()

1
2
3
训练:  model.save_weights()    
加载:  model = LeNet()
	  model.load_weights(weight_path)

① h5转pb

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
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)

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

② pb to onnx

这里需要使用netron查看pb模型的输入(input_image)，输出层(identity)名称。

转化：

1
python -m tf2onnx.convert --input ./model.pb  --inputs input_image:0 --outputs output_1:0 --output ./model.onnx --opset 12

—input: 待转化的pb模型

—inputs: 模型输入名称，有多个输入可用,分割

—outputs: 模型输出名称，有多个输出可用,分割

—opset: tf2onnx支持所有tf1.x的版本，默认情况下使用的是opset 8。当然在转换的过程中也可以指定固定的版本 —opset 12，即使用tf1.12以上的版本

四.推理测试

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
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
# onnx要求输入格式为float32，需要将双精度进行转换
img= img.astype(np.float32)
# 加载onnx
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]])

如图推理结果一致即可: