將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)