tf.keras.utils.warmstart_embedding_matrix

View source on GitHub

Warm start embedding matrix with changing vocab.

tf.keras.utils.warmstart_embedding_matrix(
    base_vocabulary,
    new_vocabulary,
    base_embeddings,
    new_embeddings_initializer='uniform'
)

This util can be used to warmstart the embedding layer matrix when vocabulary changes between previously saved checkpoint and model. Vocabulary change could mean, the size of the new vocab is different or the vocabulary is reshuffled or new vocabulary has been added to old vocabulary. If the vocabulary size changes, size of the embedding layer matrix also changes. This util remaps the old vocabulary embeddings to the new embedding layer matrix.

Example:

Here is an example that demonstrates how to use the warmstart_embedding_matrix util.

>>> import keras
>>> vocab_base = tf.convert_to_tensor(["unk", "a", "b", "c"])
>>> vocab_new = tf.convert_to_tensor(
...        ["unk", "unk", "a", "b", "c", "d", "e"])
>>> vectorized_vocab_base = np.random.rand(vocab_base.shape[0], 3)
>>> vectorized_vocab_new = np.random.rand(vocab_new.shape[0], 3)
>>> warmstarted_embedding_matrix = warmstart_embedding_matrix(
...       base_vocabulary=vocab_base,
...       new_vocabulary=vocab_new,
...       base_embeddings=vectorized_vocab_base,
...       new_embeddings_initializer=keras.initializers.Constant(
...         vectorized_vocab_new))

Here is an example that demonstrates how to get vocabulary and embedding weights from layers, use the warmstart_embedding_matrix util to remap the layer embeddings and continue with model training.

# get old and new vocabulary by using layer.get_vocabulary()
# for example assume TextVectorization layer is used
base_vocabulary = old_text_vectorization_layer.get_vocabulary()
new_vocabulary = new_text_vectorization_layer.get_vocabulary()
# get previous embedding layer weights
embedding_weights_base = model.get_layer('embedding').get_weights()[0]
warmstarted_embedding = keras.utils.warmstart_embedding_matrix(
                              base_vocabulary,
                              new_vocabulary,
                              base_embeddings=embedding_weights_base,
                              new_embeddings_initializer="uniform")
updated_embedding_variable = tf.Variable(warmstarted_embedding)

# update embedding layer weights
model.layers[1].embeddings = updated_embedding_variable
model.fit(..)
# continue with model training