问题描述:
笔者在试着用java读取深度模型时,在模型转化为byte[]传入Graph后,模型在读取输出的时候读取的维度不正确。
原因分析:
经过笔者的多次尝试,终于找到了问题的关键。问题出在h5转pb的过程中。下面是网上常见的h5转pb的python代码。使用如下代码所转化的pb文件,在使用java读取的时候就会出现输出的维度不是java想要的维度。
#*-coding:utf-8-*
"""
将keras的.h5的模型文件,转换成TensorFlow的pb文件
"""
# ==========================================================
from keras.models import load_model
import tensorflow as tf
import os.path as osp
import os
from keras import backend
#from keras.models import Sequential
def 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')
解决方案:
大家可以调试看一下h5_model.output[i]和h5_model.outputs[i]的区别,上述代码是h5_model.output[i],所得到的输出shape是一维的,而使用h5_model.outputs[i]得到的shape是二维的,正好符合java所需的维度。下面的完整的代码。
#*-coding:utf-8-*
"""
将keras的.h5的模型文件,转换成TensorFlow的pb文件
"""
# ==========================================================
from keras.models import load_model
import tensorflow as tf
import os.path as osp
import os
from keras import backend
#from keras.models import Sequential
def 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.outputs[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')
如果想学习在后端部署深度学习模型,请参考:Tomcat部署图片分类模型系列(一)