环境:
tensorflow1.15,cuda10.0,cudnn7.6.4
将keras训练好保存的.hdf5格式模型转为tensorflow的.pb模型,然后转为tensorrt支持的uff格式。
keras(.hdf5)模型转TensorFlow(.pb)
# h5_to_pb.py
from keras.models import load_model
import tensorflow as tf
import os
import os.path as osp
from keras import backend as K
# 路径参数
input_path = 'D:/pycharm/facenet/models/'
weight_file = 'resnet50.hdf5'
weight_file_path = osp.join(input_path, weight_file)
output_graph_name = weight_file[:-3] + '.pb'
# 转换函数
def h5_to_pb(h5_model, output_dir, model_name, out_prefix="output_", log_tensorboard=True):
if osp.exists(output_dir) == False:
os.mkdir(output_dir)
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))
sess = K.get_session()
from tensorflow.python.framework import graph_util, graph_io
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)
if log_tensorboard:
from tensorflow.python.tools import import_pb_to_tensorboard
import_pb_to_tensorboard.import_to_tensorboard(osp.join(output_dir, model_name), output_dir)
# 输出路径
output_dir = osp.join(os.getcwd(), "trans_model")
# 加载模型
print(weight_file_path)
h5_model = load_model(weight_file_path)
# model.load_weights
h5_to_pb(h5_model, output_dir=output_dir, model_name=output_graph_name)
print('model saved')
TensorFlow(.pb) 转TensorRT(.uff)
只需要安装好TensorRT之后通过终端运行convert-to-uff resnet50.pb即可得到resnet50.uff 。