我是在实际工程中要用到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