TensorFlow:恢复多个图形

时间:2016-10-18 02:50:30

标签: python machine-learning tensorflow

假设我们有两个TensorFlow计算图,G1G2,并保存了权重W1W2。假设我们只需构建GG1即可构建新图G2。我们如何为这个新图W1恢复W2G

举个简单的例子:

import tensorflow as tf

V1 = tf.Variable(tf.zeros([1]))
saver_1 = tf.train.Saver()
V2 = tf.Variable(tf.zeros([1]))
saver_2 = tf.train.Saver()

sess = tf.Session()
saver_1.restore(sess, 'W1')
saver_2.restore(sess, 'W2')

在此示例中,saver_1成功恢复了相应的V1,但saver_2失败并显示NotFoundError

1 个答案:

答案 0 :(得分:4)

您可以使用两个储户,每个储户只查找其中一个变量。如果您只使用tf.train.Saver(),我认为它会查找您定义的所有变量。您可以使用tf.train.Saver([v1, ...])为其提供要查找的变量列表。有关详细信息,请在此处阅读tf.train.Saver构造函数:https://www.tensorflow.org/versions/r0.11/api_docs/python/state_ops.html#Saver

这是一个简单的工作示例。假设你在一个文件中进行计算" save_vars.py"它有以下代码:

import tensorflow as tf

# Graph 1 - set v1 to have value [1.0]
g1 = tf.Graph()
with g1.as_default():
    v1 = tf.Variable(tf.zeros([1]), name="v1")
    assign1 = v1.assign(tf.constant([1.0]))
    init1 = tf.initialize_all_variables()
    save1 = tf.train.Saver()

# Graph 2 - set v2 to have value [2.0]
g2 = tf.Graph()
with g2.as_default():
    v2 = tf.Variable(tf.zeros([1]), name="v2")
    assign2 = v2.assign(tf.constant([2.0]))
    init2 = tf.initialize_all_variables()
    save2 = tf.train.Saver()

# Do the computation for graph 1 and save
sess1 = tf.Session(graph=g1)
sess1.run(init1)
print sess1.run(assign1)
save1.save(sess1, "tmp/v1.ckpt")

# Do the computation for graph 2 and save
sess2 = tf.Session(graph=g2)
sess2.run(init2)
print sess2.run(assign2)
save2.save(sess2, "tmp/v2.ckpt")

如果您确保拥有tmp目录并运行python save_vars.py,那么您将获得已保存的检查点文件。

现在,您可以使用名为" restore_vars.py"的文件进行恢复。使用以下代码:

import tensorflow as tf

# The variables v1 and v2 that we want to restore
v1 = tf.Variable(tf.zeros([1]), name="v1")
v2 = tf.Variable(tf.zeros([1]), name="v2")

# saver1 will only look for v1
saver1 = tf.train.Saver([v1])
# saver2 will only look for v2
saver2 = tf.train.Saver([v2])
with tf.Session() as sess:
    saver1.restore(sess, "tmp/v1.ckpt")
    saver2.restore(sess, "tmp/v2.ckpt")
    print sess.run(v1)
    print sess.run(v2)

当您运行python restore_vars.py时,输出应为

[1.]
[2.]

(至少在我的电脑上输出了)。如果有任何不清楚的地方,请随意发表评论。