keras模型h5转换为tensorflow/serving的pb格式模型

from keras import models
import keras.backend as K
from tensorflow.python.framework import graph_util

import tensorflow._api.v2.compat.v1 as tf
tf.disable_eager_execution()

'''
tf.saved_model.builder.SavedModelBuilder
这种模型转换的模型可以提供给TensorFlow Serving使用
多一个Variable文件
'''

def keras_model_to_tfs(model, export_path):
    signature = tf.saved_model.signature_def_utils.predict_signature_def(
        inputs={'input_x': model.input}, 
        outputs={'output_y': model.output}
    )
    builder = tf.saved_model.builder.SavedModelBuilder(export_path)
    legacy_init_op = tf.group(tf.tables_initializer(), name='legacy_init_op')
    builder.add_meta_graph_and_variables(
        sess=K.get_session(),
        tags=[tf.saved_model.tag_constants.SERVING],
        signature_def_map={
            tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature,
        },
        legacy_init_op=legacy_init_op)
    builder.save()
    print('转换完成！')


if __name__=='__main__':
    network=models.load_model('./model/mnist_cnn.h5')
    export_path='pb2'
	keras_model_to_tfs(network,export_path)

运行以上代码，就能将keras的模型文件h5转换为TensorFlow serving所需的pb格式文件。