将预训练的嵌入导入Tensorflow的嵌入特征列

时间:2019-10-07 22:52:36

标签: tensorflow gensim word2vec transfer-learning

我有一个TF估计器,在其输入层使用功能列。其中之一是EmbeddingColumn,我已经对其进行了随机初始化(默认行为)。

现在,我想对gensim中的嵌入进行预训练,并将学习到的嵌入转移到我的TF模型中。 embedding_column接受一个初始化程序参数,该参数期望使用tf.contrib.framework.load_embedding_initializer的{​​{3}}可以调用。

但是,该函数需要保存的TF检查点,而我没有,因为我在gensim中训练了嵌入。

问题是:如何将gensim词向量(为numpy数组)保存为TF检查点格式的张量,以便可以使用它初始化我的嵌入列?

1 个答案:

答案 0 :(得分:0)

想通了!这在Tensorflow 1.14.0中有效。

您首先需要将嵌入矢量转换为tf.Variable。然后使用tf.train.Saver将其保存在检查点中。

import tensorflow as tf
import numpy as np


ckpt_name = 'gensim_embeddings'
vocab_file = 'vocab.txt'
tensor_name = 'embeddings_tensor'
vocab = ['A', 'B', 'C']
embedding_vectors = np.array([
    [1,2,3],
    [4,5,6],
    [7,8,9]
], dtype=np.float32)

embeddings = tf.Variable(initial_value=embedding_vectors)
init_op = tf.global_variables_initializer()
saver = tf.train.Saver({tensor_name: embeddings})
with tf.Session() as sess:
    sess.run(init_op)
    saver.save(sess, ckpt_name)

# writing vocab file
with open(vocab_file, 'w') as f:
    f.write('\n'.join(vocab))

要使用此检查点来初始化嵌入功能列:

cat = tf.feature_column.categorical_column_with_vocabulary_file(
    key='cat', vocabulary_file=vocab_file)

embedding_initializer = tf.contrib.framework.load_embedding_initializer(
    ckpt_path=ckpt_name,
    embedding_tensor_name='embeddings_tensor',
    new_vocab_size=3,
    embedding_dim=3,
    old_vocab_file=vocab_file,
    new_vocab_file=vocab_file
)

emb = tf.feature_column.embedding_column(cat, dimension=3, initializer=embedding_initializer, trainable=False)

我们可以进行测试以确保已正确初始化:

def test_embedding(feature_column, sample):
    feature_layer = tf.keras.layers.DenseFeatures(feature_column)
    print(feature_layer(sample).numpy())

tf.enable_eager_execution()

sample = {'cat': tf.constant(['B', 'A'], dtype=tf.string)}

test_embedding(item_emb, sample)

预期的输出是:

[[4. 5. 6.]
 [1. 2. 3.]]

分别是“ B”和“ A”的嵌入。