2
$\begingroup$

I have a general question regarding TensorFlow's saver function.

The saver class allows us to save a session via:

saver.save(sess, "checkpoints.ckpt")

And allows us to restore the session:

saver.restore(sess, tf.train.latest_checkpoint("checkpoints.ckpt"))

Inside the TensorFlow documentation, there is an example code (with an added epoch loop, and restore):

# Create a saver.
saver = tf.train.Saver(...variables...)
# Launch the graph and train, saving the model every 1,000 steps.
sess = tf.Session()
saver.restore(sess, tf.train.latest_checkpoint("checkpoints.ckpt"))
for epoch in xrange(25):
    for step in xrange(1000000):
        sess.run(..training_op..)
        if step % 1000 == 0:
            # Append the step number to the checkpoint name:
            saver.save(sess, 'my-model', global_step=step)

The problem is that if we stopped the training loop at epoch=15, and execute again, then if we would start at epoch=0 again, but the model is trained up to epoch=15.

Is there a way to resume from epoch=15?

$\endgroup$
1
  • $\begingroup$ Could you link to the part of the documentation you are refering to? $\endgroup$ Aug 9, 2018 at 7:48

1 Answer 1

4
$\begingroup$

The network doesn't store its training progress with respect to training data - this is not part of its state, because at any point you could decide to change what data set to feed it. You could maybe modify it so that it knew about the training data and progress, stored in some tensor somewhere, but that would be unusual. So, in order to do this, you will need to save and make use of additional data outside of the TensorFlow framework.

Probably the simplest thing to do is add the epoch number to the filename. You are already adding the current step within the epoch, so just add in the epoch multiplied:

saver.save(sess, 'my-model', global_step=epoch*1000000+step)

When you load the file, you can parse the filename to discover what epoch and step you were on and use those as the start point for the xrange functions. To make this easier to re-start from any given checkpoint, you could use argparse to allow your script to take the name of the checkpoint file you want to use.

In brief, it might look like this:

# Near top of script
import argparse
import re

# Before main logic
parser = argparse.ArgumentParser()
parser.add_argument('checkpoint')
args = parser.parse_args()

start_epoch = 0
start_step = 0
if args.checkpoint:
    saver.restore(sess, tf.train.latest_checkpoint(args.checkpoint))
    found_num = re.search(r'\d+', args.checkpoint)
    if found_num:
        checkpoint_id = int(found_num.group(0))
        start_epoch = checkpoint_id // 1000000
        start_step = checkpoint_id % 1000000

# Change to xrange:
for epoch in xrange(start_epoch, 25):
    for step in xrange(start_step, 1000000):
        sess.run(..training_op..) # etc

    # At end of epoch loop, you need to re-set steps:
    start_step = 0

You may want to reduce the number of checkpoints you are creating - as it stands you would have 25,000 checkpoint files generated by your code.

Another option would be use a single checkpoint file, and to save and restore a Python pickle of a simple dict containing the state at the time you made the checkpoint, with a similar name.

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.