如何在slim.learning.train中的init_fn中加载两个检查点

时间:2017-09-29 07:20:04

标签: tensorflow tensorflow-slim

我想在使用slim.learning.train时加载两个检查点。例如,

init_fn = assign_from_checkpoint_fn(model_path, variables_to_restore)
slim.learning.train(train_op, log_dir, init_fn=init_fn)

问题是我只能在model_path中输入一个检查点文件。我想提出两个检查站。我认为可以有两种可能的解决方案:

  • 修改tf.contrib.framework.assign_from_checkpoint_fn中的以下assign_from_checkpoint_fn函数,以便model_path可以是检查点文件列表
  • 之前合并两个检查点。 (我没有找到任何工具)

有没有人帮助过我?

1 个答案:

答案 0 :(得分:1)

我找到了一个解决方案:我们可以使用如下的会话定义我们的init函数:

flow_init_assign_op, flow_init_feed_dict = slim.assign_from_checkpoint(
            flow_ckpt, flow_var_to_restore)

resnet_init_assign_op, resnet_init_feed_dict = 
slim.assign_from_checkpoint(
            resnet_ckpt, resnet_var_to_restore, ignore_missing_vars=True)

def init_fn(sess):
    sess.run(flow_init_assign_op, flow_init_feed_dict)
    sess.run(resnet_init_assign_op, resnet_init_feed_dict)