Keras训练的h5文件转pb文件并用Tensorflow加载

快来打我* 2022-12-12 15:03 799阅读 0赞

### 一、背景 ###

为了快速的搭建神经网络、训练模型，使用了Keras框架来搭建网络并进行训练，得到训练的h5模型文件后，需要将模型部署成服务，而pb格式的文件一般比较适合部署，pb模型文件的大小要比h5文件小一点，同时pb文件也适用于在TensorFlow Serving，所以需要把Keras保存的h5模型文件转成TensorFlow加载的pb格式来使用。同时本人也参考了网上几乎所有的模型格式转换的文章，经过一番尝试后，终于成功了，现将模型格式转换方法和分别使用tf1.x和tf2.x加载转换后的pb文件的总结如下。

### 二、h5文件转pb文件的方法 ###

首先声明一下，这里的h5文件都是用keras框架中的save()方法保存的

#### 方法一： ####

使用大佬写好的keras\_to\_tensorflow.py程序进行转化文件格式，项目地址：[https://github.com/amir-abdi/keras\_to\_tensorflow][https_github.com_amir-abdi_keras_to_tensorflow]

该作者提供了一份很好模型格式转换工具，能够满足绝大多数人的需求了。原理很简单：首先用Keras读取.h5模型文件，然后用 tensorflow的convert\_variables\_to\_constants函数将所有变量转换成常量，最后再write\_graph就是一个包含了网络以及参数值的 .pb文件了。

如果你的Keras模型是一个包含了网络结构和权重的h5文件，那么使用下面的命令就可以了：

python keras_to_tensorflow.py 
        --input_model="h5_model_path/model.h5" 
        --output_model="save_pb_model_path/model.pb"

以上命令包含两个参数，第一个是模型输入路径，第二个模型输出路径。输出路径即使你没创建好，代码也会帮你创建。建议使用绝对路径。

*注：该工具支持Tensorflow1.x版本*

#### 方法二： ####

使用下面的函数进行转换\*（注：该函数支持Tensorflow1.x版本）\*

def h5_to_pb(h5_model, output_dir, model_name, out_prefix="output_", log_tensorboard=True):
        """ .h5模型文件转换成pb模型文件 :param h5_model: .h5模型 :param output_dir: pb模型文件保存路径 :param model_name: pb模型文件名称 :param out_prefix: 根据训练，需要修改 :param log_tensorboard: 是否生成日志文件，默认为True :return: pb模型文件 """
        if not os.path.exists(output_dir):
            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 = backend.get_session()
    
        from tensorflow.python.framework import graph_util, graph_io
        # 写入pb模型文件
        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)

使用该函数前首先要加载h5模型文件，转换pb格式就只有两行代码

load_h5_model = load_model(h5_file_path, custom_objects=get_custom_objects())
    h5_to_pb(load_h5_model, output_dir=pb_model_path, model_name=pb_model_name)

如果模型有自定义层，加载时要在custom\_objects中写明，如

load_h5_model = load_model(h5_file_path, custom_objects={ 'CRF': CRF, 'crf_loss': crf_loss, 'crf_viterbi_accuracy': crf_viterbi_accuracy})

#### 方法三： ####

利用tf2.x版本框架冻结图结构将h5文件转pb文件，具体参考下面函数

def frozen_graph(h5_file_path, pb_model_path):
        """ 冻结模型，可以将训练好的.h5模型文件转成.pb文件 :param h5_file_path: h5模型文件路径 :param pb_model_path: pb模型文件保存路径 :return: """
        # 加载模型，如有自定义层请参考方法二末尾处如何加载
        model = tf.keras.models.load_model(h5_file_path, compile=False)
        model.summary()
    
        full_model = tf.function(lambda input_1: model(input_1))
        full_model = full_model.get_concrete_function(tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype))
    
        # Get frozen ConcreteFunction
        frozen_func = convert_variables_to_constants_v2(full_model)
        frozen_func.graph.as_graph_def()
    
        layers = [op.name for op in frozen_func.graph.get_operations()]
        # print("-" * 50)
        # print("Frozen model layers: ")
        # for layer in layers:
        # print(layer)
    
        # print("-" * 50)
        # print("Frozen model inputs: ")
        # print(frozen_func.inputs)
        # print("Frozen model outputs: ")
        # print(frozen_func.outputs)
    
        # Save frozen graph from frozen ConcreteFunction to hard drive
        tf.io.write_graph(graph_or_graph_def=frozen_func.graph,
                          logdir=pb_model_path,
                          name="model_name.pb",
                          as_text=False)
        print('model has been saved')

使用该函数时可能会遇到的错误：“AttributeError: module ‘tensorflow.python.framework.ops’ has no attribute ‘\_TensorLike’”

错误原因：tensorflow和keras版本不匹配

解决方法：升级keras或降级tensorflow，但必须保证tensorflow版本大于2.0

*以上方法所有代码在下面框架版本中测试通过：*

方法一和方法二：

keras==2.2.4

tensorflow-gpu==1.15.0

方法三：

keras==2.4.3

tensorflow-gpu==2.3.1

### 三、利用Tensorflow加载pb模型 ###

我们已经将h5文件转换成pb文件了，那现在就要测试一下文件是否能加载成功以及预测情况

#### （一）使用tf1.x框架加载由h5转成的pb文件 ####

加载及预测的主要代码如下

with tf.Session() as sess:
        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)
        # print all operation names
        for op in sess.graph.get_operations():
            print(op.name)
        # 输入（此处get_tensor_by_name方法中的参数为model的第一层op的name）
        input_x = sess.graph.get_tensor_by_name('import/input_1:0')
        # 输出（此处get_tensor_by_name方法中的参数为model的最后一层op的name）
        output = sess.graph.get_tensor_by_name('import/crf_1/one_hot:0')
        # 预测结果，input_data为向量化后的预测输入，注意输入shape要与模型保持一致
        ret = sess.run(output, { input_x: input_data})

#### （二）使用tf2.x框架加载由h5转成的pb文件 ####

加载及预测的主要代码如下

with tf.compat.v1.Session() as sess:
        with tf.io.gfile.GFile(pb_file_path, 'rb') as f:
            graph_def = tf.compat.v1.GraphDef()
            graph_def.ParseFromString(f.read())
            sess.graph.as_default()
            tf.compat.v1.import_graph_def(graph_def)
        # print all operation names
        # for op in sess.graph.get_operations():
        # print(op.name)
        # 输入
        input_x = sess.graph.get_tensor_by_name('import/input_1:0')
        # 输出
        output = sess.graph.get_tensor_by_name('import/model_1/crf_1/one_hot:0')
        # 预测结果
        ret = sess.run(output, { input_x: input_data})

整体逻辑与tf1.x版本一致，只是把tf1.x的代码改为tf2.x的写法。

最后我们将keras训练好的命名实体识别模型h5文件转成pb文件，通过tensorflow框架加载pb格式的模型文件，成功识别了新句子中的命名实体，结果如下图：

![在这里插入图片描述][20201007184915328.png_pic_center]

#### **总结：** 以上就是h5文件转pb文件的所有方法总结。最后，欢迎大家转载，转载请注明出处，谢谢。 ####

[https_github.com_amir-abdi_keras_to_tensorflow]: https://github.com/amir-abdi/keras_to_tensorflow
[20201007184915328.png_pic_center]: /images/20221123/79c18b7e6e2c4063852e931abad64e63.png