tensorflow 保存恢复模型

In this post we are going to talk about how to save the parameters into the disk and restore the saved parameters from the disk. The savable/restorable paramters of the network are Variables (i.e. weights and biases).

To save and restore your variables, all you need to do is to call the tf.train.Saver() at the end of you graph.

# create the graph
X = tf.placeholder(..)
Y = tf.placeholder(..)
w = tf.get_variale(..)
b = tf.get_variale(..)
...
loss = tf.losses.mean_squared_error(..)
optimizer = tf.train.AdamOptimizer(..).minimize(loss)
...

saver = tf.train.Saver() 

In the train mode, in the session we will initialize the variables and run our network. At the end of training, we will save the variables using saver.save():

# TRAIN
with tf.Session() as sess:
    sess.run(tf.globale_variables_initializer())
    # train our model
    for step in range(steps):
        sess.run(optimizer)
        ...
    saved_path = saver.save(sess, './my-model', global_step=step)
 

This will create 3 files (data, index, meta) with a suffix of the step you saved your model.

In the test mode, in the session we will restore the variables using saver.restore() and validate or test our model.

# TEST
with tf.Session() as sess:
    saver.restore(sess, './my-model')
    ...