how to load checkpoint file and continue training

Question

When I was training my graph, I found that I forgot to add dropout in my graph. But I have already trained a long time and get some checkpoints. So is it possible for me to load the checkpoints and add a dropout and then continue training? My code is like this now:

# create a graph 
vgg_fcn = fcn8_vgg_ours.FCN8VGG()
with tf.name_scope("content_vgg"):
    vgg_fcn.build(batch_images, train = True, debug=True)
labels = tf.placeholder("int32", [None, HEIGHT, WIDTH])
# do something
...
#####
init_glb = tf.global_variables_initializer()
init_loc = tf.local_variables_initializer()
sess.run(init_glb)
sess.run(init_loc)
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
ckpt_dir = "./checkpoints"
if not os.path.exists(ckpt_dir):
    os.makedirs(ckpt_dir)
ckpt = tf.train.get_checkpoint_state(ckpt_dir)
start = 0
if ckpt and ckpt.model_checkpoint_path:
    start = int(ckpt.model_checkpoint_path.split("-")[1])
    print("start by epoch: %d"%(start))
    saver = tf.train.Saver()
    saver.restore(sess, ckpt.model_checkpoint_path)
last_save_epoch = start
# continue training

So if I have changed the structure of FCN8VGG(add some dropout layer), then will it use the meta file to replace the graph I have just created? If it will, how could I change the structure to continue training without training from scratch again?

Answer 1

Here's a simple example of initializing a new model using variables from another model's checkpoint. Note that things are much simpler if you can just pass a variable_scope to init_from_checkpoint, but here I'm assuming that the original model was not designed with restoring in mind.

First define a simple model with some variables, and do some training:

import tensorflow as tf

def first_model():
  with tf.Graph().as_default():
    fake_input = tf.constant([[1., 2., 3., 4.],
                              [5., 6., 7., 8.]])
    layer_one_output = tf.contrib.layers.fully_connected(
        inputs=fake_input, num_outputs=5, activation_fn=None)
    layer_two_output = tf.contrib.layers.fully_connected(
        inputs=layer_one_output, num_outputs=1, activation_fn=None)
    target = tf.constant([[10.], [-3.]])
    loss = tf.reduce_sum((layer_two_output - target) ** 2)
    train_op = tf.train.AdamOptimizer(0.01).minimize(loss)
    init_op = tf.global_variables_initializer()
    saver = tf.train.Saver()
    with tf.Session() as session:
      session.run(init_op)
      for i in range(1000):
        _, evaled_loss = session.run([train_op, loss])
        if i % 100 == 0:
          print(i, evaled_loss)
      saver.save(session, './first_model_checkpoint')

Running first_model(), training looks fine and we get a first_model_checkpoint written:

0 109.432
100 0.0812649
200 8.97705e-07
300 9.64064e-11
400 9.09495e-13
500 0.0
600 0.0
700 0.0
800 0.0
900 0.0

Next, we can define a completely new model in a different graph, and initialize the variables that it shares with first_model from that checkpoint:

def second_model():
  previous_variables = [
      var_name for var_name, _
      in tf.contrib.framework.list_variables('./first_model_checkpoint')]
  with tf.Graph().as_default():
    fake_input = tf.constant([[1., 2., 3., 4.],
                              [5., 6., 7., 8.]])
    layer_one_output = tf.contrib.layers.fully_connected(
        inputs=fake_input, num_outputs=5, activation_fn=None)
    # Add a batch_norm layer, which creates some new variables. Replacing this
    # with tf.identity should verify that the model one variables are faithfully
    # restored (i.e. the loss should be the same as at the end of model_one
    # training).
    batch_norm_output = tf.contrib.layers.batch_norm(layer_one_output)
    layer_two_output = tf.contrib.layers.fully_connected(
        inputs=batch_norm_output, num_outputs=1, activation_fn=None)
    target = tf.constant([[10.], [-3.]])
    loss = tf.reduce_sum((layer_two_output - target) ** 2)
    train_op = tf.train.AdamOptimizer(0.01).minimize(loss)
    # We're done defining variables, now work on initializers. First figure out
    # which variables in the first model checkpoint map to variables in this
    # model.
    restore_map = {variable.op.name:variable for variable in tf.global_variables()
                   if variable.op.name in previous_variables}
    # Set initializers for first_model variables to restore them from the
    # first_model checkpoint
    tf.contrib.framework.init_from_checkpoint(
        './first_model_checkpoint', restore_map)
    # For new variables, global_variables_initializer will initialize them
    # normally. For variables in restore_map, they will be initialized from the
    # checkpoint.
    init_op = tf.global_variables_initializer()
    saver = tf.train.Saver()
    with tf.Session() as session:
      session.run(init_op)
      for i in range(10):
        _, evaled_loss = session.run([train_op, loss])
        print(i, evaled_loss)
      saver.save(session, './second_model_checkpoint')

In this case, previous_variables looks like:

['beta1_power', 'beta2_power', 'fully_connected/biases', 'fully_connected/biases/Adam', 'fully_connected/biases/Adam_1', 'fully_connected/weights', 'fully_connected/weights/Adam', 'fully_connected/weights/Adam_1', 'fully_connected_1/biases', 'fully_connected_1/biases/Adam', 'fully_connected_1/biases/Adam_1', 'fully_connected_1/weights', 'fully_connected_1/weights/Adam', 'fully_connected_1/weights/Adam_1']

Note that since we haven't used any variable scopes, the naming is dependent on the order layers are defined. If names change, you need to manually construct the restore_map.

If we run second_model, the loss jumps up initially because the batch_norm layer hasn't been trained:

0 38.5976
1 36.4033
2 33.3588
3 29.8555
4 26.169
5 22.5185
6 19.0838
7 16.0096
8 13.4035
9 11.3298

However, replacing batch_norm with tf.identity verifies that the previously trained variables have been restored.