问如何使检查点在tf.train优化器中存储时刻和其他相关变量

EN

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

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