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

背景：目前keras框架使用简单，很容易上手，深得广大算法工程师的喜爱，但是当部署到客户端时，可能会出现各种各样的bug，甚至不支持使用keras，本文来解决的是将keras的h5模型转换为客户端常用的tensorflow的pb模型并使用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 = '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 tf
from 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_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')

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

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

下边是两个函数介绍：

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

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

来源：https://blog.csdn.net/u010159842/article/details/84481478