SparseOperationKit Embedding Lookup

sparse_operation_kit.distributed_variable.Variable(*args, **kwargs)

Abbreviated as sok.Variable.

This is a helper function to generate model-parallel variable. There are two use cases:

Distributed Variable:

import numpy as np
import tensorflow as tf
import horovod.tensorflow as hvd
import sparse_operation_kit as sok

hvd.init()
gpus = tf.config.experimental.list_physical_devices("GPU")
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)
if gpus:
    tf.config.experimental.set_visible_devices(gpus[hvd.local_rank()], "GPU")  # nopep8

sok.init()

# If there are 2 GPUs in total, the shape on GPU0 will be [2, 3] and the shape
# on GPU1 will be [2, 3]
v = sok.Variable(np.arange(4 * 3).reshape(4, 3), dtype=tf.float32)

# GPU0 output: [[0, 1, 2]
#               [6, 7, 8]]
# GPU1 output: [[3, 4,  5]
#                9, 10, 11]
print(v)

Localized Variable:

import numpy as np
import tensorflow as tf
import horovod.tensorflow as hvd
import sparse_operation_kit as sok

hvd.init()
gpus = tf.config.experimental.list_physical_devices("GPU")
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)
if gpus:
    tf.config.experimental.set_visible_devices(gpus[hvd.local_rank()], "GPU")  # nopep8

sok.init()

# If there are 2 GPUs in total, the shape on GPU0 will be [5, 3] and the shape
# on GPU1 will be [0, 3]
v = sok.Variable(
    np.arange(5 * 3).reshape(5, 3), dtype=tf.float32, mode="localized:0"
)
print(v.shape)

As shown in the two examples above, when you need to store different parts of a variable on different GPUs (that is, allocating a model-parallel variable), this function can help you allocate the required memory on each GPU.

Parameters:

• args – compatible with tf.Variable.

• kwargs – compatible with tf.Variable.

• mode (string) – a string to specify which model-parallel mode to use. Default value is “distributed”, which stands for the Distributed Variable that mentioned above. Another option is “localized:#”, which stands for Localized Variable, where # indicates which GPU you want to put this variable on. See the explanation above for specific examples.

Returns:

variable – a tf.Variable that represents a part of the model-parallel variable.

Return type:

tf.Variable

sparse_operation_kit.lookup.lookup_sparse(params, sp_ids, sp_weights=None, combiners=None)

Abbreviated as sok.lookup_sparse.

Perform fused sparse lookup on the given embedding params. This function is similar to the tf.nn.embedding_lookup_sparse, but with two differences:

• It can do distributed lookup.

• It can accept multiple params and multiple sp_ids to do fused lookup at once, which brings performance benefits.

Parameters:

• params (list, tuple) – a list or tuple of trainable sok.Variable.

• sp_ids (list, tuple) – a list or tuple of tf.SparseTensor or tf.RaggedTensor.

• sp_weights (list tuple,optional) – a list or tuple of tf.SparseTensor or tf.RaggedTensor(float / double weights). if don’t specify , indicate all weights should be taken to be 1.

• combiners (list, tuple,optional) – a list or tuple of string to specify the combiner of each lookup,for now only suupport “mean” “sum”. if don’t specify , indicate all elements(numbero of elements is same with number of sok.Variables) in combiners will should be set to be mean.

Returns:

emb_vec – a list of tf.Tensor(the results of lookup).

Return type:

list

Example

import numpy as np
import tensorflow as tf
import horovod.tensorflow as hvd
import sparse_operation_kit as sok

hvd.init()
gpus = tf.config.experimental.list_physical_devices("GPU")
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)
if gpus:
    tf.config.experimental.set_visible_devices(gpus[hvd.local_rank()], "GPU")

sok.init()

v1 = sok.Variable(np.arange(17 * 3).reshape(17, 3), dtype=tf.float32)
v2 = sok.Variable(np.arange(7 * 5).reshape(7, 5), dtype=tf.float32)

indices1 = tf.SparseTensor(
    indices=[[0, 0], [0, 1], [1, 0], [1, 1], [1, 2]],
    values=[1, 1, 3, 4, 5],
    dense_shape=[2, 3])
)
indices2 = tf.SparseTensor(
    indices=[[0, 0], [1, 0], [1, 1]],
    values=[1, 2, 3],
    dense_shape=[2, 2]
)

embeddings = sok.lookup_sparse(
    [v1, v2], [indices1, indices2], combiners=["sum", "sum"]
)
print(embeddings[0])
print(embeddings[1])