如何在tf.train优化器中创建Checkpoint存储时刻和其他相关变量

Question

I encountered a problem when my code stopped for some reason on my machine, so I had to restart my code and continue the training process by loading the latest checkpoint file.

I found that the performance is not consistent before and after the checkpoint that I loaded and the performance dropped a lot.

So, since my code uses tf.train.AdamOptimizer , I guess that the checkpoint doesn't store the moment vectors and the gradients in the previous steps, and when I load the checkpoint the moment vectors are initialized as zeros.

Am I correct?

Is there any method that can help store relevant vectors for the Adamopotimizer in the checkpoints so that if my machine is down again, restarting from the latest checkpoint will not influence anything?

Thanks!

Answer 1

Out of curiosity, I checked if it is true and everything seems to be working just fine: all variables are presented in the checkpoint and restored properly. See for yourself:

import tensorflow as tf
import sys
import numpy as np
from tensorflow.python.tools import inspect_checkpoint as inch


ckpt_path = "./tmp/model.ckpt"
shape = (2, 2)

def _print_all():
  for v in tf.all_variables():
    print('%20s' % v.name, v.eval())

def _model():
    a = tf.placeholder(tf.float32, shape)
    with tf.variable_scope("foo", reuse=tf.AUTO_REUSE):
      x = tf.get_variable('x', shape)

    loss = tf.matmul(a, tf.layers.batch_normalization(x))
    step = tf.train.AdamOptimizer(0.00001).minimize(loss)
    return a, step

def train():
    a, step = _model()
    saver = tf.train.Saver()

    with tf.Session() as sess:
      sess.run(tf.global_variables_initializer())
      for i in range(10):
        _ = sess.run(step, feed_dict= {a:np.random.rand(*shape)})

      _print_all()
      print(saver.save(sess, ckpt_path))
      _print_all()


def check():
    a, step = _model()
    saver = tf.train.Saver()

    with tf.Session() as sess:
      sess.run(tf.global_variables_initializer())
      _print_all()
      saver.restore(sess, ckpt_path)
      _print_all()


def checkpoint_list_vars(chpnt):
  """
  Given path to a checkpoint list all variables available in the checkpoint
  """
  from tensorflow.contrib.framework.python.framework import checkpoint_utils
  var_list = checkpoint_utils.list_variables(chpnt)
#   for v in var_list: print(v, var_val(v[0]))
#   for v in var_list: print(v)
  var_val('')

  return var_list

def var_val(name):
    inch.print_tensors_in_checkpoint_file(ckpt_path, name, True)

if 'restore' in sys.argv:
    check()
elif 'checkpnt' in sys.argv:
    checkpoint_list_vars(ckpt_path)
else:
    train()

Store it as test.py and run

>> python test.py
>> python test.py checkpnt
>> python test.py restore