TensorFlow学习笔记：保存和读取模型

TensorFlow 更新频率实在太快，从 1.0 版本正式发布后，很多 API 接口就发生了改变。今天用 TF 训练了一个 CNN 模型，结果在保存模型的时候居然遇到各种问题。Google 搜出来的答案也是莫衷一是，有些回答对 1.0 版本的已经不适用了。后来实在没办法，就翻了墙去官网看了下，结果分分钟就搞定了～囧～。

这篇文章内容不多，主要讲讲 TF v1.0 版本中保存和读取模型的最简单用法，其实就是对官网教程的简要翻译摘抄。

保存和恢复

在 TensorFlow 中，保存和恢复模型最简单的方法就是使用 tf.train.Saver 类。这个类会将变量的保存和恢复操作添加到 TF 的图（graph）中。

Checkpoint 文件

TF 将变量保存在二进制文件中，这个文件包含一个从变量名到 tensor 值的映射。当我们创建一个 Saver 对象的时候，我们可以指定 checkpoint 文件中的变量名。默认会使用变量的 Variable.name 属性。

这一段读起来比较生涩难懂，具体看下面的例子。

保存变量

可以通过创建 Saver 来管理模型内的所有变量。

# Create some variables.
v1 = tf.Variable(..., name="v1")
v2 = tf.Variable(..., name="v2")
...
# Add an op to initialize the variables.
init_op = tf.global_variables_initializer()

# Add ops to save and restore all the variables.
saver = tf.train.Saver()

# Later, launch the model, initialize the variables, do some work, save the
# variables to disk.
with tf.Session() as sess:
  sess.run(init_op)
  # Do some work with the model.
  ..
  # Save the variables to disk.
  save_path = saver.save(sess, "/tmp/model.ckpt")
  print("Model saved in file: %s" % save_path)

恢复变量

可以通过同一个 Saver 对象（指定相同的保存路径）来恢复变量。这种情况下，我们不需要事先初始化变量（即无需调用 tf.global_variables_initializer()）

# Create some variables.
v1 = tf.Variable(..., name="v1")
v2 = tf.Variable(..., name="v2")
...
# Add ops to save and restore all the variables.
saver = tf.train.Saver()

# Later, launch the model, use the saver to restore variables from disk, and
# do some work with the model.
with tf.Session() as sess:
  # Restore variables from disk.
  saver.restore(sess, "/tmp/model.ckpt")
  print("Model restored.")
  # Do some work with the model
  ...

例子

下面用我自己的例子解释一下。

首先，我们先定义一个图模型（只截选出变量部分）：

    graph = tf.Graph()

    with graph.as_default():
        # Input data
        # ....省略代码若干

        # Variables
        layer1_weights = tf.Variable(tf.truncated_normal(
            [patch_size, patch_size, image_channels, depth], stddev=0.1), name="layer1_weights")
        layer1_biases = tf.Variable(tf.zeros([depth]), name="layer1_biases")

        layer2_weights = tf.Variable(tf.truncated_normal(
            [image_size // 4 * image_size // 4 * depth, num_hidden], stddev=0.1, name="layer2_weights")
        )
        layer2_biases = tf.Variable(tf.constant(1.0, shape=[num_hidden]), name="layer2_biases")

        layer3_weights = tf.Variable(tf.truncated_normal(
            [num_hidden, num_labels], stddev=0.1, name="layer3_weights"),
        )
        layer3_biases = tf.Variable(tf.constant(1.0, shape=[num_labels]), name="layer3_biases")

        def model(data):
            #....省略代码若干
            return tf.matmul(fc1, layer3_weights) + layer3_biases

        # Training computation
        #....省略代码若干

        # Optimizer
        optimizer = tf.train.GradientDescentOptimizer(0.05).minimize(loss)

这个模型里的变量其实只有三个网络层的参数：layer1_weights，layer1_biases，layer2_weights，layer2_biases，layer3_weights，layer3_biases。

然后就是启动会话进行训练：

    with tf.Session(graph=graph) as session:
        saver = tf.train.Saver()

        if loading_model:
            saver.restore(session, model_folder + "/" + model_file)
            print("Model restored")
        else:
            tf.global_variables_initializer().run()
            print("Initialized")

        for step in range(num_steps):
            # ....省略训练模型的代码

        print('Test accuracy: %.1f%%' % accuracy(test_prediction.eval(), test_labels))
        save_path = saver.save(session, model_folder + "/" + model_file)
        print("Model saved in file: ", save_path)

这段代码是本文的关键，我们先通过 tf.train.Saver() 构造一个 Saver 对象，注意，这一步要在 Session 启动之后执行，否则会抛异常 ValueError("No variables to save")，至少 v1.0 是这样。

通过 Saver，我们可以在模型训练完之后，将参数保存下来。Saver 保存数据的方法十分简单，只要将 session 和文件路径传入 save 函数即可：saver.save(session, model_folder + "/" + model_file)。

如果我们一开始想载入本地的模型文件，而不是让 TF 自动初始化训练，则可以通过 Saver 的 restore 函数读取模型文件，文件路径需要和之前保存的文件路径一致。注意，如果是通过这种方式初始化变量，则不能再调用 tf.global_variables_initializer() 函数。之后，训练或预测的代码不需要改变，TensorFlow 会自动根据模型文件，将你的模型参数初始化。

当然啦，以上都是最基础的用法，只是简单地将所有参数保存下来。更高级的用法，之后如果使用到再继续总结。

参考

• TensorFlow官方教程