How to get the global_step when restoring checkpoints in Tensorflow?

Viewed 26204

I'm saving my session state like so:

self._saver = tf.saver()
self._saver.save(self._session, '/network', global_step=self._time)

When I later restore I want to get the value of the global_step for the checkpoint I restore from. This is in order to set some hyper parameters from it.

The hacky way to do this would be to run through and parse the file names in the checkpoint directory. But surly there has to be a better, built in way to do this?

8 Answers

You can use the global_step variable to keep track of steps, but if in your code, you are initializing or assigning this value to another step variable, it may not be consistent.

For instance, you define your global_step using:

global_step = tf.Variable(0, name='global_step', trainable=False)

Assign to your training operation:

train_op = optimizer.minimize(loss, global_step=global_step)

Save in your checkpoint:

saver.save(sess, checkpoint_path, global_step=global_step)

And restore from your checkpoint:

saver.restore(sess, checkpoint_path) 

the value of global_step is restored as well but if you are assigning it to another variable, say step, then you must do something like:

step = global_step.eval(session=sess)

The variable step, contains the last saved global_step in the checkpoint.

It will be nice to also define the global_step from graph than as zero variable (as earlier defined):

global_step = tf.train.get_or_create_global_step()

This will get your last global_step if exist or create one if not.

TL;DR

As tensorflow variable (will be evaluated in the session)

global_step = tf.train.get_or_create_global_step()
# use global_step variable to calculate your hyperparameter 
# this variable will be evaluated later in the session
saver = tf.train.Saver()
with tf.Session() as sess:
    # restore all variables from checkpoint
    saver.restore(sess, checkpoint_path)
    # than init table and local variables and start training/evaluation ...

Or: As numpy integer (without any session):

reader = tf.train.NewCheckpointReader(absolute_checkpoint_path)
global_step = reader.get_tensor('global_step')


Long Answer

There are at least two ways retrieving the global from a checkpoint. As tensorflow variable or as numpy integer. Parsing the filename will not work, if the global_step was not provided as a parameter in the save method of the Saver. For pretrained models see the remark at the end of the answer.

As Tensorflow variable

If you need the global_step variable to calculate some hyperparameters you can just use tf.train.get_or_create_global_step(). This will return a tensorflow variable. Because the variable will be evaluated later in the session you can only use tensorflow operations to calculate your hyperparameters. So e.g.: max(global_step, 100) will not work. You have to use tensorflow equivalent tf.maximum(global_step, 100) that can be evaluated later in the session.

Within the session you can initialize the global step variable with a checkpoint using saver.restore(sess, checkpoint_path)

global_step = tf.train.get_or_create_global_step()
# use global_step variable to calculate your hyperparameter 
# this variable will be evaluated later in the session
hyper_parameter = tf.maximum(global_step, 100) 
saver = tf.train.Saver()
with tf.Session() as sess:
    # restore all variables from checkpoint
    saver.restore(sess, checkpoint_path)
    # than init table and local variables and start training/evaluation ...

    # for verification you can print the global step and your hyper parameter
    print(sess.run([global_step, hyper_parameter]))

Or: As numpy integer (without session)

If you need the global step variable as scalar without starting a session you can also read this variable directly from your checkpoint file(s). You just need a NewCheckpointReader. Because of a bug in older tensorflow versions you should convert the path of the checkpoint file to an absolute path. With the reader you can get all the tensors of the model as numpy variables. The name of the global step variable is a constant string tf.GraphKeys.GLOBAL_STEP defined as 'global_step'.

absolute_checkpoint_path = os.path.abspath(checkpoint_path)
reader = tf.train.NewCheckpointReader(absolute_checkpoint_path)
global_step = reader.get_tensor(tf.GraphKeys.GLOBAL_STEP)

Remark to pretrained models: In most pretrained models that are available online the global step is reset to zero. So, these models can be used to initialize the model parameters for finetuning without overwrite the global step.

just note my solution on global step saving and restore.

Save:

global_step = tf.Variable(0, trainable=False, name='global_step')
saver.save(sess, model_path + model_name, global_step=_global_step)

Restore:

if os.path.exists(model_path):
    saver.restore(sess, tf.train.latest_checkpoint(model_path))
    print("Model restore finished, current globle step: %d" % global_step.eval())

The reason that a variable is not restored as expected is most likely due to the fact that it was created after your tf.Saver() object was created.

The place where you create the tf.Saver() object matters when you don't explicitly specify a var_list, or specify None for var_list. The expected behavior for many programmers is that all variables in the graph are saved when the save() method is called, but this is not the case, and it should perhaps be documented as such. A snapshot of all variables in the graph is saved at the time of object creation.

Unless you're having any performance issues, it's safest to create the saver object right when you decide to save your progress. Otherwise, make sure to create the saver object after you create all your variables.

Also, the global_step that is passed to saver.save(sess, save_path, global_step=global_step) is merely a counter used for creating the filename and has nothing to do with whether it will be restored as a global_step variable. This is a parameter misnomer IMO since if you're saving your progress at the end of each epoch, it's probably best to pass your epoch number for this parameter.

Related