tensorflow的ckpt及pb模型持久化方式及转化详解

1. ckpt和pb模型的持久化方式

在tensorflow中，我们通常使用ckpt和pb文件来保存和加载模型。ckpt文件是tensorflow的Checkpoint文件，保存了模型的权重和偏置等参数信息，并且还包含了计算图的结构信息。而pb文件是tensorflow的Protocol Buffer文件，它是一种可序列化的数据结构，可以同时保存模型的计算图和参数信息，可以方便地部署到其他平台上。

1.1 ckpt模型持久化

ckpt模型持久化是指将tensorflow模型的参数信息保存为ckpt文件的过程。在tensorflow中，我们可以使用tf.train.Saver类来实现模型的保存和加载。

import tensorflow as tf

# 定义模型
x = tf.placeholder(tf.float32, shape=(None,))
y = tf.placeholder(tf.float32, shape=(None,))
z = tf.add(x, y)
# 创建Saver对象
saver = tf.train.Saver()
# 保存模型
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    saver.save(sess, 'model.ckpt')

在上面的代码中，我们首先定义了一个简单的模型，然后创建了一个Saver对象。最后，在会话中调用saver.save()方法保存模型，将参数信息保存在model.ckpt文件中。

1.2 pb模型持久化

pb模型持久化是指将tensorflow模型的计算图和参数信息保存为pb文件的过程。在tensorflow中，我们可以使用tf.saved_model.builder.SavedModelBuilder类来实现模型的保存和加载。

import tensorflow as tf

# 定义模型
x = tf.placeholder(tf.float32, shape=(None,))
y = tf.placeholder(tf.float32, shape=(None,))
z = tf.add(x, y)
# 创建SavedModelBuilder对象
builder = tf.saved_model.builder.SavedModelBuilder('model.pb')
# 保存模型
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    builder.add_meta_graph_and_variables(sess, ['my_model'])
    builder.save()

在上面的代码中，我们定义了一个简单的模型，然后创建了一个SavedModelBuilder对象，传入保存pb文件的路径。最后，在会话中调用builder.add_meta_graph_and_variables()方法保存模型，将计算图和参数信息保存在model.pb文件中。

2. ckpt模型的转化

在实际应用中，我们有时候需要将ckpt模型转化为pb模型，以便在其他平台上部署和使用模型。tensorflow提供了一个工具 called freeze_graph.py，可以将ckpt模型转化为pb模型。

2.1 安装freeze_graph.py工具

freeze_graph.py工具是tensorflow的一个Python脚本，需要通过pip安装。在命令行中执行以下命令：

pip install tensorflow

2.2 使用freeze_graph.py工具

在命令行中执行以下命令：

python freeze_graph.py --input_checkpoint=model.ckpt --output_graph=model.pb --output_node_names=Add

其中，--input_checkpoint参数指定了ckpt模型的路径，--output_graph参数指定了要保存的pb模型的路径，--output_node_names参数指定了要转化的计算图中的输出节点的名称，这里我们假设输出节点的名称为"Add"。

3. pb模型的加载和使用

在其他平台上部署和使用pb模型时，我们需要加载pb模型，并对输入数据进行预测。

import tensorflow as tf

# 加载pb模型
graph = tf.Graph()
with tf.gfile.FastGFile('model.pb', 'rb') as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
    with graph.as_default():
        tf.import_graph_def(graph_def, name='')
# 使用pb模型进行预测
with tf.Session(graph=graph) as sess:
    inputs = graph.get_tensor_by_name('input:0')
    output = graph.get_tensor_by_name('Add:0')
    result = sess.run(output, feed_dict={inputs: [[1, 2]]})
    print(result)

在上面的代码中，我们首先加载了pb模型，使用tf.gfile.FastGFile()方法读取pb文件，然后使用tf.GraphDef()类解析读取的pb文件数据，创建一个新的图。最后，在会话中调用sess.run()方法进行预测，通过feed_dict传入输入数据，并获取输出结果。

4. 总结

本文详细介绍了tensorflow中ckpt和pb模型的持久化方式及转化方法。通过ckpt模型的保存和加载，我们可以方便地保存和恢复tensorflow模型的参数信息；通过pb模型的保存和加载，我们可以将tensorflow模型的计算图和参数信息一起保存，方便在其他平台上部署和使用模型。同时，我们还介绍了如何使用freeze_graph.py工具将ckpt模型转化为pb模型，并给出了pb模型的加载和使用示例代码。