machine-learning - 如何加载检查点文件并使用略有不同的图形结构继续训练

Question

当我训练图表时，我发现我忘记在图表中添加 dropout。但我已经训练了很长时间并得到了一些检查点。那么我是否可以加载检查点并添加 dropout，然后继续训练？我的代码现在是这样的:

# 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

那么如果我改变了FCN8VGG的结构(添加一些dropout层)，那么它会使用元文件来替换我刚刚创建的图吗？如果可以的话，如何改变结构继续训练而不需要再次从头开始训练？

Answer 1

这是一个使用另一个模型检查点的变量初始化新模型的简单示例。请注意，如果您只需将 variable_scope 传递给 init_from_checkpoint，事情就会简单得多，但这里我假设原始模型在设计时并未考虑到恢复。

首先定义一个带有一些变量的简单模型，并进行一些训练:

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')

运行first_model()，训练看起来不错，我们得到了一个first_model_checkpoint:

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

接下来，我们可以在不同的图中定义一个全新的模型，并从该检查点初始化它与first_model共享的变量:

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')

在这种情况下，previous_variables 看起来像:

['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']

请注意，由于我们没有使用任何变量范围，因此命名取决于定义层的顺序。如果名称发生变化，您需要手动构建restore_map。

如果我们运行 second_model，损失最初会上升，因为 batch_norm 层尚未经过训练:

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

但是，将 batch_norm 替换为 tf.identity 可以验证之前训练的变量是否已恢复。