.h5 To Pb

以下代碼實現了把.h5文件轉換成pb的過程

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 = '/home/simon/ASTAR Validation/Code4GE_23092019/Astar trained model'
weight_file = 'trained_model_run1.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(),"high_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')

ps：如果生產的.pb文件不能用，可能是tensorlow 版本的問題