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。一个好的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

在Estimators

中

对于Estimators，我们无法直接访问会话。初始化嵌入的方法可以是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)

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

Answer 2

奥利维尔·莫恩德罗特（Olivier Moindrot）已经给出了好的方法，我只是添加了我发现的另一种方法。为了避免出现2GB的限制错误，您应该真正关心所使用的操作。张量流中的图可序列化为protobuf（谷歌设计的交换格式）。在检查操作的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中初始化嵌入层

2 个答案:

在Estimators