Tensorflow中的tf.constant_initializer非常缓慢

4

尝试使用预训练的100维word2vec嵌入来训练LSTM。

@staticmethod
def load_embeddings(pre_trained_embeddings_path, word_embed_size):
    embd = []
    import time
    start_time = time.time()
    cnt = 4
    with codecs.open(pre_trained_embeddings_path, mode="r", encoding='utf-8') as f:
        for line in f.readlines():
            values = line.strip().split(' ')
            embd.append(values[1:])
            cnt += 1
            if cnt % 100000 == 0:
                print("word-vectors loaded: %d" % cnt)

    embedding, vocab_size, embed_dim = embd, len(embd), len(embd[0])

    load_end_time = time.time()
    print("word vectors loaded from and start initialising, cnt: %d, time taken: %d secs " % (vocab_size, load_end_time - start_time))

    embedding_init = tf.constant_initializer(embedding, dtype=tf.float16)
    src_word_embedding = tf.get_variable(shape=[vocab_size, embed_dim], initializer=embedding_init, trainable=False, name='word_embedding', dtype=tf.float16)

    print("word-vectors loaded and initialised, cnt: %d, time taken: %d secs" % (vocab_size, time.time() - load_end_time))

    return src_word_embedding

这个方法执行后的输出结果如下:
word vectors loaded from and start initialising, cnt: 2419080, time taken: 74 secs
word-vectors loaded and initialised, cnt: 2419080, time taken: 1647 secs

系统信息:tensorflow 1.1.0,tcmalloc,python 3.6,ubuntu 14.04

初始化需要半个小时似乎非常慢,这是正常现象还是有问题?有什么想法或问题吗?

更新:使用@sirfz提供的嵌入方法使得加载嵌入变得非常快速,初始化完成只需85秒


这似乎是一个未解决的问题。请参考GPU上的布尔运算非常缓慢 - YLJ
是的,与 float32 类似的时间。 - jknair
2个回答

0

将大常量加载到图中不仅速度较慢,还会泄漏大量内存。我曾经遇到过类似的问题不久前报告,对我来说最好的解决方法是:

# placeholder for loading your saved embeddings
embedding_init = tf.placeholder(tf.float16, shape=[vocab_size, embed_dim])
src_word_embedding = tf.get_variable(initializer=embedding_init, trainable=False, name='word_embedding', dtype=tf.float16)

# run initialization with the value of embeddings placeholder
session.run(tf.global_variables_initializer(), feed_dict={embedding_init: embedding})

在保存检查点时,这会保存嵌入吗?在恢复时,这会加载回原始的嵌入吗? - jknair
这个简化了很多时间,太棒了!初始化完成,用时85秒 - jknair
我第一次尝试时也有同样的反应 ;) 至于检查点,我认为它会正常工作,因为它只会将变量序列化。个人而言,我更喜欢使用 numpy.save 来转储变量,因为这样更快。 - sirfz

0

我不知道这是否是一种预期的行为,但我可以解释为什么在一个小例子上它会变慢:

import tensorflow as tf

x = [[0, 1], [2, 3]]
a = tf.constant(x, name='a')
b = tf.Variable(x, name='b')
c = a + b

with tf.Session() as sess:
    writer = tf.summary.FileWriter('logs', sess.graph)
    writer.close()

当您初始化常量时,此常量的值将添加到图形中。如果您打开图形,可以通过单击a值来查看它。

enter image description here

在我的情况下,它是一个2x2的矩阵,但看起来在你的情况下是一个2M x ?的矩阵,非常巨大。所以我认为这就是执行缓慢的原因。
尝试将其初始化为一个变量并在那里提供您的嵌入。

确切地说,这是将大张量序列化为图中的节点,这是减速(最终导致更大的内存占用)的主要来源。使用我回答中的解决方法可以避免这个昂贵的序列化步骤。 - sirfz

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接