Tensorflow 使用pb文件保存(恢复)模型计算图和参数

1. 背景介绍

Tensorflow 是由 Google 开源的一个用于构建和训练神经网络的深度学习框架，其广泛应用于图像、语音、自然语言处理等领域。在 Tensorflow 中，通常使用计算图来表示模型，计算图由节点和边组成，节点表示操作，边表示数据流向。Tensorflow 支持将计算图和参数保存为 pb 文件，方便模型的复用和部署。

2. pb 文件介绍

pb 文件是 Protocol Buffer 格式的文件，类似于 XML 或 JSON 格式的文件，可以用于序列化和反序列化数据。在 Tensorflow 中，pb 文件不仅能保存计算图，还能保存模型参数，其通常包含以下文件：

模型计算图的 pb 文件，通常以 .pb 作为扩展名；

模型参数的 ckpt 文件，通常包含多个文件，以 .ckpt 开头，例如 .ckpt.data-00000-of-00001，.ckpt.index，.ckpt.meta 等。

在模型的使用中，通常只需要使用 .pb 文件即可，因为通过 .pb 文件可以加载完整的计算图和参数。如果只有 ckpt 文件，需要先将其转换为 pb 文件，方法见下一节。

3. 保存模型为 pb 文件

在 Tensorflow 中，可以使用 tf.train.Saver 类来保存模型参数。下面是一个简单的示例：

import tensorflow as tf

# 定义计算图
x = tf.Variable(initial_value=3.0, name='x')
y = tf.square(x)
z = tf.sqrt(y)
# 创建 tf.train.Saver 对象
saver = tf.train.Saver()
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    # 进行一些训练操作
    # ...
    # 保存模型为 pb 文件
    saver.save(sess, 'model.pb')

上述代码中，定义了一个简单的计算图，包含一个变量 x、一个平方操作 y 和一个开方操作 z。创建了一个 tf.train.Saver 对象 saver，并在 Session 中调用 saver.save 方法保存模型为 pb 文件（'model.pb'）。

需要注意的是，如果计算图中包含占位符等需要用户手动 feed 的变量，则需要在保存模型时指定输入占位符，方法如下：

...

# 定义计算图
x = tf.placeholder(tf.float32, name='x')
y = tf.square(x)
z = tf.sqrt(y)
# 创建 tf.train.Saver 对象，定义输入占位符
saver = tf.train.Saver({'x': x})
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    # 进行一些训练操作
    # ...
    # 保存模型为 pb 文件，指定输入占位符
    saver.save(sess, 'model.pb', global_step=global_step)

4. 加载 pb 文件

加载 pb 文件可以使用 tf.gfile.GFile 和 tf.GraphDef 函数。下面是一个简单的示例：

import tensorflow as tf

with tf.gfile.GFile('model.pb', 'rb') as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
with tf.Graph().as_default() as graph:
    tf.import_graph_def(graph_def)
with tf.Session(graph=graph) as sess:
    # 使用模型进行推理
    x = graph.get_tensor_by_name('x:0')
    y = graph.get_tensor_by_name('power:0')
    z = graph.get_tensor_by_name('sqrt:0')
    result = sess.run(z, feed_dict={x: 2.0})

上述代码中，首先使用 tf.gfile.GFile 函数打开 pb 文件并读取其中的数据，然后使用 tf.GraphDef 函数将数据解析为一个计算图定义。接着在新的计算图中导入计算图定义，这样就可以使用新的计算图加载模型了。

需要注意的是，如果计算图中包含占位符等需要用户手动 feed 的变量，在使用模型进行推理时需要先获取相应的输入和输出 tensors，方法如下：

...

# 加载 pb 文件
with tf.gfile.GFile('model.pb', 'rb') as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
with tf.Graph().as_default() as graph:
    tf.import_graph_def(graph_def)
# 获取输入和输出 tensors
x = graph.get_tensor_by_name('x:0')
y = graph.get_tensor_by_name('power:0')
z = graph.get_tensor_by_name('sqrt:0')
with tf.Session(graph=graph) as sess:
    # 使用模型进行推理
    result = sess.run(z, feed_dict={x: 2.0})

5. 将 ckpt 文件转换为 pb 文件

如果只有 ckpt 文件，需要先将其转换为 pb 文件。可以使用 tf.train.import_meta_graph 函数导入 ckpt 文件，然后使用 tf.train.write_graph 函数将计算图保存为 pb 文件。下面是一个简单的示例：

import tensorflow as tf

# 导入 ckpt 文件
saver = tf.train.import_meta_graph('model.ckpt.meta')
# 将计算图保存为 pb 文件
with tf.Session() as sess:
    saver.restore(sess, 'model.ckpt')
    graph_def = tf.get_default_graph().as_graph_def()
    tf.train.write_graph(graph_def, '.', 'model.pb', as_text=False)

上述代码中，首先使用 tf.train.import_meta_graph 函数导入 ckpt 文件，并在 Session 中调用 saver.restore 方法恢复模型参数。然后使用 tf.get_default_graph 函数获取计算图，并使用 tf.train.write_graph 函数将计算图保存为 pb 文件。

需要注意的是，如果计算图中包含占位符等需要用户手动 feed 的变量，在保存 pb 文件时需要指定输入占位符，方法如下：

...

# 获取输入占位符
x = tf.get_default_graph().get_tensor_by_name('x:0')
# 将计算图保存为 pb 文件，指定输入占位符
tf.train.write_graph(graph_def, '.', 'model.pb', as_text=False, input_map={'x:0': x})

6. 总结

本文介绍了如何将 Tensorflow 中的计算图和参数保存为 pb 文件，并在另一个计算图中加载 pb 文件。同时还介绍了如何将 ckpt 文件转换为 pb 文件。通过 pb 文件，可以方便地保存、加载和部署 Tensorflow 模型。