tensorflow加载大型词嵌入模型（＞2G）

最新推荐文章于 2024-07-19 23:55:01 发布

人工智能小a

最新推荐文章于 2024-07-19 23:55:01 发布

阅读量406

点赞数

CC 4.0 BY-SA版权

文章标签： tensorflow python 深度学习

本文链接：https://blog.youkuaiyun.com/a_123zxy/article/details/125835350

本文介绍了在TensorFlow中加载大型词嵌入模型的两种方法。当文件大小适中时，可以直接通过`tf.get_variable`和`tf.nn.embedding_lookup`加载。对于超过2GB的大文件，可以使用`tf.Variable`结合`placeholder`和`assign`操作，但在运行时可能会降低网络训练速度。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

tensorflow加载词嵌入模型

方法一

当词嵌入文件较小的时候，直接加载：

embedding = tf.get_variable(initializer=word_embeddings, name='label_embedding')
word_ebd = tf.nn.embedding_lookup(embedding, word_id)

方法二

当词嵌入文件较大的时候（>2G），加载方法如下：

embedding_weights = tf.Variable(tf.constant(0.0, shape=[embedding_vocab_size, EMBEDDING_DIM]), trainable=False, name='embedding_weights')
embedding_placeholder = tf.placeholder(tf.float32, [embedding_vocab_size, EMBEDDING_DIM])
embedding = embedding_weight.assign(embedding_placeholder)
word_ebd = tf.nn.embedding_lookup(embedding, word_id)

sess = tf.Session()
word_ebd = sess.run(word_ebd, feed_dict={embedding_placeholder: word_embeddings, word_id: word_id})