概述
该方法只适用于通过save_model保存的keras模型文件(.h5),不适用于通过save_weights保存的keras模型文件,直接上代码。
import keras
from keras.models import load_model
import tensorflow as tf
import os.path as osp
import os
from keras import backend as K
# 转换函数
def h5_to_pb(h5_model, output_dir, new_name, out_prefix="output_", log_tensorboard=True):
if not osp.exists(output_dir):
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)
# h5模型路径
h5_model_name = 'model.h5'
log_path = osp.join(os.getcwd(), h5_model_name)
# pb文件输出路径
output_dir = osp.join(os.getcwd(), "pb_model")
# 会自行创建"pb_model"文件夹
# 输出模型名称
output_graph_name = h5_model [:-3] + '.pb'
# 模型冻结方式
K.set_learning_phase(0)
# 0 testing mode, 1 training mode
# 加载模型
h5_model = load_model(log_path )
h5_to_pb(h5_model, output_dir=output_dir, new_name=output_graph_name)
print('model saved')
最后
以上就是阳光铃铛为你收集整理的keras模型h5文件转pb文件(通过save_model保存的h5模型文件)的全部内容,希望文章能够帮你解决keras模型h5文件转pb文件(通过save_model保存的h5模型文件)所遇到的程序开发问题。
如果觉得靠谱客网站的内容还不错,欢迎将靠谱客网站推荐给程序员好友。
本图文内容来源于网友提供,作为学习参考使用,或来自网络收集整理,版权属于原作者所有。
发表评论 取消回复