Importing pre-trained embeddings into Tensorflow's Embedding Feature Column - tensorflow

Question

I have a TF Estimator that uses Feature Columns at its input layer. One of these is and EmbeddingColumn which I have been initializing randomly (the default behaviour).
Now I would like to pre-train my embeddings in gensim and transfer the learned embeddings into my TF model. The embedding_column accepts an initializer argument which expects a callable that can be created using tf.contrib.framework.load_embedding_initializer.
However, that function expects a saved TF checkpoint, which I don't have, because I trained my embeddings in gensim.
The question is: how do I save gensim word vectors (which are numpy arrays) as a tensor in the TF checkpoint format so that I can use that to initialize my embedding column?

Answer 1

Figured it out! This worked in Tensorflow 1.14.0.
You first need to turn the embedding vectors into a tf.Variable. Then use tf.train.Saver to save it in a checkpoint.
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))
To use this checkpoint to initialize an embedding feature column:
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)
And we can test to make sure it has been initialized properly:
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)
The output, as expected, is:
[[4. 5. 6.]
[1. 2. 3.]]
Which are the embeddings for 'B' and 'A' respectively.