向经过训练的 Tensorflow 估算器添加层

Adding a layer to a trained Tensorflow estimator

我有一个使用 Estimator API 训练的 Tensorflow 模型。现在我想加载这个模型,向它添加一个新层,并只训练新层(即冻结所有其他参数)。

最好的方法是什么?

我设法从检查点加载模型并添加新层,但我不清楚如何将示例作为输入提供并执行训练。具体来说,我找不到输入的占位符。

我找到了一个简单的方法来做到这一点,并将其张贴在这里以防有人发现它有用:

  1. 加载 Estimator 检查点。
  2. 在新名称范围下创建占位符和模型图的副本。
  3. 提取两个作用域下的所有可训练变量。
  4. 为每个变量创建 assign 个操作。

代码:

    # Load the trained model from checkpoint.
    new_saver = tf.train.import_meta_graph('{}.meta'.format(config.ckpt_fullpath))
    new_saver.restore(sess, config.ckpt_fullpath)

    # Create new graph with a placeholder for input.
    new_model_scope = 'new_scope'
    trained_model_scope = 'old_scope' # this should be taken from the original model function of the estimator.
    with tf.name_scope(new_model_scope):
        model = Model(config)
        input_tensor = tf.placeholder(tf.float32, 
                                      [None, config.img_size[0], config.img_size[1], 3])
        model.build_model(input_tensor)

    # Initialize the new graph variables with trained parameters.
    trained_params = [t for t in tf.trainable_variables()
                      if t.name.startswith(trained_model_scope)]
    trained_params = sorted(trained_params, key=lambda v: v.name)
    new_params = [t for t in tf.trainable_variables() 
                  if t.name.startswith(new_model_scope)]
    new_params = sorted(new_params, key=lambda v: v.name)

    update_ops = []
    for trained_v, new_v in zip(trained_params, new_params):
        op = new_v.assign(trained_v)
        update_ops.append(op)
    sess.run(update_ops)