将keras训练好的h5模型转换成TensorFlow serving环境调用的.pb文件的方法

最近做人脸识别项目，以前一直是在服务器上训练好模型，用TensorFlow serving 保存成pb文件调用。后来发现 Keras的h5模型文件也可以转换成Pb文件以供TensorFlow serving调用。实在是大快人心啊！
下面看看如何将Keras的h5转换成Pb模型文件吧：

方法一

Keras模型保存时使用 save() 函数保存，转换代码如下：

from keras.models import load_model
import tensorflow as tf
import os, sys
from keras import backend as K
from tensorflow.python.framework import graph_util, graph_io

def h5_to_pb(h5_weight_path, output_dir, out_prefix = "output_", log_tensorboard = True):

    if not os.path.exists(output_dir):
        os.mkdir(output_dir)
        
    h5_model = load_model(h5_weight_path)
    
    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))
    
    model_name = os.path.splitext(os.path.split(h5_weight_path)[-1])[0] + '.pb'    
    
    sess = K.get_session()
    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__':
    if len(sys.argv) == 3:
        # usage: python3 h5_to_pb.py h5_weight_path output_dir
        h5_to_pb(h5_weight_path=sys.argv[1], output_dir=sys.argv[2])

方法二

Keras模型保存时使用 save_weigths() 函数，转换代码如下：

from keras.models import load_model
import tensorflow as tf
import os, sys
from keras import backend as K
from tensorflow.python.framework import graph_util, graph_io

def h5_to_pb(h5_weight_path, output_dir, out_prefix = "output_", log_tensorboard = True):

    if not os.path.exists(output_dir):
        os.mkdir(output_dir)
        
    h5_model = build_model()
    h5_model.load_weights(h5_weight_path)
    
    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))
    
    model_name = os.path.splitext(os.path.split(h5_weight_path)[-1])[0] + '.pb'    
    
    sess = K.get_session()
    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)
        
        
def build_model():
    
    inputs = Input(shape=(784, ), name='input_img')
    x = Dense(64, activation='relu')(inputs)
    x = Dense(64, activation='relu')(x)
    y = Dense(10, activation='softmax')(x)
    h5_model = Model(inputs=inputs, outputs=y)
    
    return h5_model
    
        
if __name__ == '__main__':
    if len(sys.argv) == 3:
        # usage: python3 h5_to_pb.py h5_weight_path output_dir
        h5_to_pb(h5_weight_path=sys.argv[1], output_dir=sys.argv[2])

注意：
1、函数build_model()根据自己的实际模型做修改即可
2、“ inputs = Input(shape=(784, ), name='input_img') ” 中name='input_img'最好是自己命名，在pb使用时会用到