TensorFlow 检查点保存和读取

TensorFlow checkpoint save and read

我有一个基于 TensorFlow 的神经网络和一组变量。

训练函数是这样的:

def train(load = True, step)
    """
    Defining the neural network is skipped here
    """

    train_step = tf.train.AdamOptimizer(1e-4).minimize(mse)
    # Saver
    saver = tf.train.Saver()

    if not load:
        # Initalizing variables
        sess.run(tf.initialize_all_variables())
    else:
        saver.restore(sess, 'Variables/map.ckpt')
        print 'Model Restored!'

    # Perform stochastic gradient descent
    for i in xrange(step):
        train_step.run(feed_dict = {x: train, y_: label})

    # Save model
    save_path = saver.save(sess, 'Variables/map.ckpt')
    print 'Model saved in file: ', save_path
    print 'Training Done!'

我是这样调用训练函数的:

# First train
train(False, 1)
# Following train
for i in xrange(10):
    train(True, 10)

我进行这种训练是因为我需要为我的模型提供不同的数据集。但是,如果我以这种方式调用训练函数,TensorFlow 将生成错误消息,指示它无法从文件中读取保存的模型。

经过一些实验,我发现这是因为检查点保存很慢。在文件写入磁盘之前,下一个训练函数将开始读取,从而产生错误。

我曾尝试使用 time.sleep() 函数在每次调用之间进行一些延迟,但没有成功。

有人知道如何解决这种 write/read 错误吗?非常感谢!

您的代码中存在一个微妙的问题:每次您调用 train() 函数时,都会将更多节点添加到同一个 TensorFlow 图,用于所有模型变量和神经网络的其余部分。这意味着每次构造 tf.train.Saver() 时,它都会包含之前调用 train() 的所有变量。每次重新创建模型时,都会创建带有额外 _N 后缀的变量,以赋予它们唯一的名称:

  1. 用变量var_avar_b.
  2. 构建的Saver
  3. 用变量 var_avar_bvar_a_1var_b_1.
  4. 构建的 Saver
  5. 用变量var_avar_bvar_a_1var_b_1var_a_2var_b_2
  6. 构造的Saver
  7. 等等

tf.train.Saver 的默认行为是将每个变量与相应操作的名称相关联。这意味着 var_a_1 不会从 var_a 初始化,因为它们最终有不同的名称。

解决方案是每次调用 train() 时创建一个新图。修复它的最简单方法是更改​​主程序,为每次调用 train() 创建一个新图形,如下所示:

# First train
with tf.Graph().as_default():
    train(False, 1)

# Following train
for i in xrange(10):
    with tf.Graph().as_default():
        train(True, 10)

...或者,等效地,您可以将 with 块移动到 train() 函数中。