SparseOperationKit Embedding Saver

Currently, you must explicitly call this saver to dump or load trainable parameters for the trainable variables in those embedding layers. And other trainable variables are still managed by the DL framework.

class sparse_operation_kit.saver.Saver.Saver[source]

Abbreviated as sok.Saver().

dump_to_file(embedding_variable, filepath)[source]

This function is used to save the specified embedding variables to host file.

When multiple CPU processes is used, this function must be called within each CPU processes.

Parameters

  • embedding_variable (sok.EmbeddingVariable, tf.DistributedVariable) – The variable from embedding layer which needs to be dumped to file.

  • filepath (string) – The directory where the parameters will be dumped to.

Returns

status – If this op executed successfully, then ‘OK’ will be returned.

Return type

tf.Tensor

restore_from_file(embedding_variable, filepath)[source]

This function is used to restore dumped parameters to the specified embedding variable.

When multiple CPU processes is used, this function must be called within each CPU processes.

Parameters

  • embedding_variable (sok.EmbeddingVariable, tf.DistributedVariable) – The embedding variable which needs to be restored from file.

  • filepath (string) – The directory where the parameters will be restored from.

Returns

status – If this op executed successfully, then ‘OK’ will be returned.

Return type

tf.Tensor

load_embedding_values(embedding_variable, tensors)[source]

This function is used to assign embedding_variable’s value with tf.Tensors.

When multiple CPU processes is used, this function must be called within each CPU processes.

Parameters

  • embedding_variable (sok.EmbeddingVariable, tf.DistributedVariable) – Which embedding_variable’s value will be assigned.

  • tensors (tf.Tensor, list of tf.Tensor, tuple of tf.Tensor) – Each tf.Tensor must be 2-rank and the shape must be [None, embedding_vec_size], where the embedding_vec_size must be equal to that of embedding_variable’s. All tf.Tensors make up to a big tensor, which just like they are stacked. For example: [tf.Tensor(shape=(bs_0, embedding_vec_size)), tf.Tensor(shape=(bs_1, embedding_vec_size)), tf.Tensor(shape=(bs_2, embedding_vec_size))] will be treated as tf.Tensor(shape=(bs_0 + bs_1 + bs_2, embedding_vec_size)).

Returns

status – If this op executed successfully, then ‘OK’ will be returned.

Return type

tf.Tensor