如何在Tensorflow中恢复模型

时间:2018-04-11 13:31:54

标签: python tensorflow tensorflow-slim

首先,我使用tf.contrib.gan训练了一个模型,我能够训练模型。

tf.contrib.gan.gan_train(
        train_ops,
        hooks=(
                [tf.train.StopAtStepHook(num_steps=FLAGS.max_number_of_steps),
                 tf.train.LoggingTensorHook([status_message], every_n_iter=10)] +
                sync_hooks),
        logdir=FLAGS.train_log_dir,
        master=FLAGS.master,
        is_chief=FLAGS.task == 0,
        config=conf
    )

然后我想评估模型。尝试以下列方式恢复检查点时,

with tf.name_scope('inputs'):
  real_images, one_hot_labels, _, num_classes = data_provider.provide_data(
    FLAGS.batch_size, FLAGS.dataset_dir)
  logits, end_points_des, feature, net = dcgan.discriminator(real_images)

  variables_to_restore = slim.get_model_variables()
  restorer = tf.train.Saver(variables_to_restore)

  with tf.Session() as sess:
          ckpt = tf.train.get_checkpoint_state(FLAGS.checkpoint_dir)
          restorer.restore(sess, ckpt.model_checkpoint_path)

我得到了这个例外:

      2018-04-11 20:05:03.304089: W tensorflow/core/framework/op_kernel.cc:1192] Not found: Key Discriminator/fully_connected_layer2/weights not found in checkpoint
      2018-04-11 20:05:03.304280: W tensorflow/core/framework/op_kernel.cc:1192] Not found: Key Discriminator/conv0/BatchNorm/Discriminator/conv0/BatchNorm/moving_mean/local_step not found in checkpoint
      2018-04-11 20:05:03.304484: W tensorflow/core/framework/op_kernel.cc:1192] Not found: Key Discriminator/conv0/BatchNorm/beta not found in checkpoint
      2018-04-11 20:05:03.305197: W tensorflow/core/framework/op_kernel.cc:1192] Not found: Key Discriminator/fully_connected_layer2/biases not found in checkpoint

我正在使用TF 1.7rc1

1 个答案:

答案 0 :(得分:0)

实际上,生成的图表中存在问题。这些是我为解决这个问题所做的步骤。

第1步:  使用以下代码

打印checkpoit文件中的所有变量
from tensorflow.python.tools.inspect_checkpoint import print_tensors_in_checkpoint_file
print_tensors_in_checkpoint_file(file_name, '')

步骤2:然后我注意到每个键都包含第一个范围(' Discriminator')的复制,但是当我尝试加载模型时,它不包含该部分。因此,我必须以下列方式删除该附加部分,

  def name_in_checkpoint(var):
      if "Discriminator/" in var.op.name:
         return var.op.name.replace("Discriminator/", "Discriminator/Discriminator/")

  logits, end_points_des, feature, net = dcgan.discriminator(real_images) 
  variables_to_restore = slim.get_model_variables()
  variables_to_restore = {name_in_checkpoint(var): var for var in variables_to_restore}
  restorer = tf.train.Saver(variables_to_restore)

步骤3:然后我能够成功加载模型,如下所示。

ckpt = tf.train.get_checkpoint_state(FLAGS.checkpoint_dir)
restorer.restore(sess, ckpt.model_checkpoint_path)

希望这会帮助可能遇到同样问题的人。