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

Question

我正在尝试在 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.

进行此类初始化，将会很有帮助

Answer 1

如果您将 initializer 参数指定为 tf.get_variable()，初始值 GLOVE_MATRIX 将存储在图表中并超过 2Gb。很好地解释了如何加载嵌入。

这是第一个示例，其中我们使用了初始化程序，图 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 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)

使用内置脚手架的一个好处是它只会在您第一次调用 train_input_fn 时初始化嵌入。对于以后的调用，它不会再运行 init_fn。

Answer 2

Olivier Moindrot 已经给出了很好的方法，我只是添加我发现的另一种方法。为了避免 2GB 限制错误，您应该真正关心您使用的操作。 tensorflow 中的 Graph 可序列化为 protobuf（google 设计的一种交换格式）。当检查attr_value 超过2GB 时创建操作时会触发2GB 错误，因此我们必须避免将大值打包到这个attr_value 操作字段中。我只是提供了另一个函数tf.variable.load，这个函数不会在Graph中添加额外的东西示例代码如下：

graph2 = tf.Graph()
size = 1000
toy_embedding = np.random.randn(size, size)  # you can create your own np.ndarray

with graph2.as_default() as g2:
    emb2 = tf.Variable(initial_value=tf.ones(shape=(size,size)), name='emb2')
    sess = tf.Session()
    init_op = tf.global_variables_initializer()
    sess.run(init_op)
    emb2.load(toy_embedding, session = sess)
    out = sess.run(emb2)
    print (out)

graph2.as_graph_def().ByteSize()

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

How to initialize embeddings layer within Estimator API

tensorflow

tensorflow-estimator

在估算器中