tensorflow:如何保存/恢复训练模型

时间:2017-08-31 00:17:53

标签: tensorflow

这是一个如何保存和恢复训练模型的示例。 希望这对初学者有所帮助。

生成具有relu激活功能的1个隐藏层神经网络。 (听说relu被证明比sigmoid好得多,特别是对于有大量隐藏层的神经网络。)

训练数据显然是异或。

训练并保存“tf_train_save.py”

import tensorflow as tf
import numpy as np

x = np.matrix([[0, 0], [0, 1], [1, 0], [1, 1]])
y = np.matrix([[0], [1], [1], [0]])

n_batch = x.shape[0]
n_input = x.shape[1]
n_hidden = 5
n_classes = y.shape[1]

X = tf.placeholder(tf.float32, [None, n_input], name="X")
Y = tf.placeholder(tf.float32, [None, n_classes], name="Y")

w_h = tf.Variable(tf.random_normal([n_input, n_hidden], stddev=0.01), tf.float32, name="w_h")
w_o = tf.Variable(tf.random_normal([n_hidden, n_classes], stddev=0.01), tf.float32, name="w_o")

l_h = tf.nn.relu(tf.matmul(X, w_h))
hypo = tf.nn.relu(tf.matmul(l_h, w_o), name="output")

cost = tf.reduce_mean(tf.square(Y-hypo))
train = tf.train.GradientDescentOptimizer(0.1).minimize(cost)

init = tf.global_variables_initializer()

with tf.Session() as sess:
    sess.run(init)

    for epoch in range(1000):
        for i in range(4):
            sess.run(train, feed_dict = {X:x[i,:], Y:y[i,:]})

    result = sess.run([hypo, tf.floor(hypo+0.5)], feed_dict={X:x})

    print(*result[0])
    print(*result[1])

    output_graph_def = tf.graph_util.convert_variables_to_constants(sess, sess.graph_def, ["output"])
    tf.train.write_graph(output_graph_def, "./logs/mp_logs", "test.pb", False)

加载“tf_load.py”

import tensorflow as tf
from tensorflow.python.platform import gfile
import numpy as np

x = np.matrix([[0, 0], [0, 1], [1, 0], [1, 1]])
y = np.matrix([[0], [1], [1], [0]])

with gfile.FastGFile("./logs/mp_logs/test.pb",'rb') as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
    tf.import_graph_def(graph_def, name='')

with tf.Session() as sess:
    X = sess.graph.get_tensor_by_name("X:0")
    print(X)
    output = sess.graph.get_tensor_by_name("output:0")
    print(output)

    tf.global_variables_initializer().run()

    result = sess.run([output, tf.floor(output+0.5)], feed_dict={X:x})

    print(*result[0])
    print(*result[1])

会有更简单的方法吗?

1 个答案:

答案 0 :(得分:1)

您正在使用convert_variables_to_constants,因此您在训练方面确实很好。对于路人而言,该API出现在v1.0中(如果我在跟踪API之后没有弄错)。

在负载方面,我认为最小代码是一个命令更短。鉴于您已将所有变量都转换为常量,因此在还原时没有要初始化的变量。这一行:

tf.global_variables_initializer().run()

什么都不做。来自v1.3的docs

  

但是,如果var_list为空,则该函数仍会返回可以运行的Op。 Op只是没有效果。

加载脚本没有全局变量,并且由于tf.global_variables_initializer()等同于tf.variables_initializer(tf.global_variables()),因此操作是无操作。