tensorflow - 如何在Estimator API中初始化嵌入层

我正在尝试在tensorflow模型中使用现有的嵌入，嵌入的大小大于2Gb，这使我的原始尝试失败了：

embedding_var = tf.get_variable(
        "embeddings",
        shape=GLOVE_MATRIX.shape,
        initializer=tf.constant_initializer(np.array(GLOVE_MATRIX))
)

这给了我这个错误：

 Cannot create a tensor proto whose content is larger than 2GB.

我正在使用基于Estimator API的AWS SageMaker，并且会话中图形的实际运行发生在幕后，因此我不确定如何为嵌入而初始化一些占位符。如果有人能够使用EstimatorAPI共享如何进行此类初始化的方式，这将很有帮助。

最佳答案

如果将initializer参数指定为tf.get_variable()，则初始值GLOVE_MATRIX将存储在图形中并超过2Gb。好的answer解释了一般如何加载嵌入内容。



这是第一个示例，其中我们使用初始化器，并且图def大约为4Mb，因为它在其中存储了(1000, 1000)矩阵。

size = 1000
initial_value = np.random.randn(size, size)
x = tf.get_variable("x", [size, size], initializer=tf.constant_initializer(initial_value))

sess = tf.Session()
sess.run(x.initializer)

assert np.allclose(sess.run(x), initial_value)

graph = tf.get_default_graph()
print(graph.as_graph_def().ByteSize())  # should be 4000394

这是一个更好的版本，我们不存储它：

size = 1000
initial_value = np.random.randn(size, size)
x = tf.get_variable("x", [size, size])

sess = tf.Session()
sess.run(x.initializer, {x.initial_value: initial_value})

assert np.allclose(sess.run(x), initial_value)

graph = tf.get_default_graph()
print(graph.as_graph_def().ByteSize())  # should be 1203

在估算器中

对于估算者，我们没有直接访问该会话的权限。初始化嵌入的一种方法是使用tf.train.Scaffold。您可以为其传递参数init_fn，在其中初始化嵌入变量，而无需将实际值保存在图形def中。

def model_fn(features, labels, mode):
    size = 10
    initial_value = np.random.randn(size, size).astype(np.float32)
    x = tf.get_variable("x", [size, size])

    def init_fn(scaffold, sess):
        sess.run(x.initializer, {x.initial_value: initial_value})
    scaffold = tf.train.Scaffold(init_fn=init_fn)

    loss = ...
    train_op = ...

    return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op, scaffold=scaffold)

使用内置Scaffold的一个好处是，只有在您首次调用train_input_fn时，它才会初始化嵌入。对于以后的呼叫，它将不再运行init_fn。

关于tensorflow - 如何在Estimator API中初始化嵌入层，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/48217599/