将keras的h5模型转换为tensorflow的pb模型操作

时间:2021-05-22

背景:目前keras框架使用简单,很容易上手,深得广大算法工程师的喜爱,但是当部署到客户端时,可能会出现各种各样的bug,甚至不支持使用keras,本文来解决的是将keras的h5模型转换为客户端常用的tensorflow的pb模型并使用tensorflow加载pb模型。

h5_to_pb.py from keras.models import load_modelimport tensorflow as tfimport os import os.path as ospfrom keras import backend as K#路径参数input_path = 'input path'weight_file = 'weight.h5'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")#加载模型h5_model = load_model(weight_file_path)h5_to_pb(h5_model,output_dir = output_dir,model_name = output_graph_name)print('model saved')

将转换成的pb模型进行加载

load_pb.py import tensorflow as tffrom tensorflow.python.platform import gfile def load_pb(pb_file_path): sess = tf.Session() with gfile.FastGFile(pb_file_path, 'rb') as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) sess.graph.as_default() tf.import_graph_def(graph_def, name='') print(sess.run('b:0')) #输入 input_x = sess.graph.get_tensor_by_name('x:0') input_y = sess.graph.get_tensor_by_name('y:0') #输出 op = sess.graph.get_tensor_by_name('op_to_store:0') #预测结果 ret = sess.run(op, {input_x: 3, input_y: 4}) print(ret)

补充知识:h5模型转化为pb模型,代码及排坑

我是在实际工程中要用到tensorflow训练的pb模型,但是训练的代码是用keras写的,所以生成keras特定的h5模型,所以用到了h5_to_pb.py函数。

附上h5_to_pb.py(python3)

#*-coding:utf-8-*"""将keras的.h5的模型文件,转换成TensorFlow的pb文件"""# ==========================================================from keras.models import load_modelimport tensorflow as tfimport os.path as ospimport osfrom keras import backend#from keras.models import Sequentialdef h5_to_pb(h5_model, output_dir, model_name, out_prefix="output_", log_tensorboard=True): """.h5模型文件转换成pb模型文件 Argument: h5_model: str .h5模型文件 output_dir: str pb模型文件保存路径 model_name: str pb模型文件名称 out_prefix: str 根据训练,需要修改 log_tensorboard: bool 是否生成日志文件 Return: pb模型文件 """ if os.path.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 = backend.get_session() from tensorflow.python.framework import graph_util, graph_io # 写入pb模型文件 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(os.path.join(output_dir, model_name), output_dir)if __name__ == '__main__': # .h模型文件路径参数 input_path = 'D:/CSP' weight_file = 'xingren.h5' weight_file_path = os.path.join(input_path, weight_file) output_graph_name = weight_file[:-3] + '.pb' # pb模型文件输出输出路径 output_dir = osp.join(os.getcwd(),"trans_model") #model.save(xingren.h5) # 加载模型 #h5_model = Sequential() h5_model = load_model(weight_file_path) #h5_model.save(weight_file_path) #h5_model.save('xingren.h5') h5_to_pb(h5_model, output_dir=output_dir, model_name=output_graph_name) print ('Finished')

在运行的时候遇到了下面问题:

原因:我们训练模型的时候用save_weights函数保存模型,但是这个函数只保存了权重文件,并没有又保存模型的参数。要把save_weights改为save。

下边是两个函数介绍:

save()保存的模型结果,它既保持了模型的图结构,又保存了模型的参数。

save_weights()保存的模型结果,它只保存了模型的参数,但并没有保存模型的图结构

以上这篇将keras的h5模型转换为tensorflow的pb模型操作就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持。

声明:本页内容来源网络,仅供用户参考;我单位不保证亦不表示资料全面及准确无误,也不保证亦不表示这些资料为最新信息,如因任何原因,本网内容或者用户因倚赖本网内容造成任何损失或损害,我单位将不会负任何法律责任。如涉及版权问题,请提交至online#300.cn邮箱联系删除。

相关文章