如何使用预训练变量初始化 TensorFlow 估计器的模型函数中的权重？答案

【问题标题】：How to initialize weights in the model function of a TensorFlow estimator, using pre-trained variables?如何使用预训练变量初始化 TensorFlow 估计器的模型函数中的权重？
【发布时间】：2018-05-31 06:39:54
【问题描述】：

我想使用 TensorFlow 的 tf.Estimator API 包装我的模型，所以我有以下model_fn：

    def model_fn(features, labels, mode, params):
        # Load vocabulary
        n_vocab = params['n_vocab']

        if mode != tf.estimator.ModeKeys.PREDICT:
            labels = tf.reshape(labels, (-1, 1))

        embedding = tf.Variable(tf.random_uniform((tf.cast(n_vocab, tf.int32),
                                                  tf.cast(params['embed_space'], tf.int32)), 0, 1), name='embedding')
        embedding_layer = tf.nn.embedding_lookup(embedding, features[INPUT_TENSOR_NAME], name='embedding_layer')

        ...

它初始化嵌入变量并执行一些其他操作。现在，我不想使用tf.random_uniform 初始化这些权重，而是使用我之前训练过的已保存模型的权重，并冻结它们，使它们不可训练。我怎样才能做到这一点？

【问题讨论】：

标签： tensorflow tensorflow-estimator

【解决方案1】：

您可以使用tf.estimator.WarmStartSettings。 Estimator 通过添加一个名为 warm_start_from 的参数提供了一种简单的方法来创建它。

创建估算器时，应传递保存的模型目录：

estimator = tf.estimator.Estimator(
    model_fn, model_dir, warm_start_from=SAVED_MODEL_DIR)

请注意，您可能需要已保存模型的解冻版本。

【讨论】：

甚至在使用`tf.estimator.LinearRegressor`开始训练之前，我们需要为我们定义的模型初始化权重，tensorflow中的哪个API可以做到这一点？你能指点我吗？可能是 tf 自己在做的事情，但是如何以及从哪里开始？感谢您回答在哪里加载预训练的权重，但不清楚 tf 在第一次运行中从哪里开始。