TensorFlow1.5训练模型的保存与加载

           TensorFlow1.5训练模型的保存与加载

  

模型保存：

import numpy as np
import tensorflow as tf
param = np.array([[1], [2], [3]])
graph = tf.get_default_graph()
#名称也很重要，方便后续的载入
w = tf.Variable([0], dtype=tf.float32,name="w")
x = tf.placeholder(tf.float32, [3, 1],name="x")
costFunction = x[0][0] * w ** 2 + x[1][0] * w + x[2][0]
#关键点，将模型的引用保存，后续载入时可直接进行run
graph.add_to_collection(name="cost", value=costFunction)
train = tf.train.GradientDescentOptimizer(0.01).minimize(costFunction)
init = tf.global_variables_initializer()
saver = tf.train.Saver()

with tf.Session() as sesson:
    sesson.run(init)
    print(sesson.run(w))
    for i in range(10):
        sesson.run(train, feed_dict={x : param})
        #方便在观察损失函数的变化情况
        print(sesson.run(costFunction, feed_dict={x : param}))
    saver.save(sesson, "E:/data/model/model_test")
    print(sesson.run(w))

模型加载：

import tensorflow as tf
import numpy as np
with tf.Session() as session:
    saver = tf.train.import_meta_graph('E:/data/model/model_test.meta')
    saver.restore(session, tf.train.latest_checkpoint('E:/data/model/'))
    #print(session.run("w:0"))
    graph = tf.get_default_graph()
    #w = graph.get_tensor_by_name('w:0')
    #根据名称加载placeholder
    x = graph.get_tensor_by_name('x:0')
    param = np.array([[1], [2], [3]])
    #根据名称获取保存的损失函数引用
    cost = graph.get_collection("cost")
    print(session.run(cost,feed_dict={x : param}))