SparseOperationKit Utilities

sparse_operation_kit.dump_load.dump(path, dump_vars, optimizer=None)

Abbreviated as sok.dump.

Dump the embedding tables into one folder , will generate key and value file per table.

Table name prefix is same as variable name in tensorflow.

Now is only support SGD,Adamax,Adadelta,Adagrad,Ftrl,Adam optimizers.

Note: for Adam optimizers , it functions in the same way as LazyAdam in SOK, rather than TensorFlow’s Adam

Parameters:

• path (string) – weight file folder

• dump_vars (List,Tuple,SOK Variable) – Can be a single or list of sok.Variable and sok.DynamicVariable

• optimizer (SOK.OptimizerWrapper,optional,default is None) – when model train , need to dump optimizer state,input sok.OptimizerWrapper

Return type:

None

Example

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

v = sok.DynamicVariable(dimension=3, var_type="hbm", initializer="13")

optimizer = tf.keras.optimizers.SGD(learning_rate=1.0)
optimizer = sok.OptimizerWrapper(optimizer)
path = "your weight path , weight is dumped by sok.dump"

indices = tf.convert_to_tensor([0, 1, 2**40], dtype=tf.int64)

with tf.GradientTape() as tape:
    embedding = tf.nn.embedding_lookup(v, indices)
    print("embedding:", embedding)
    loss = tf.reduce_sum(embedding)

grads = tape.gradient(loss, [v])
optimizer.apply_gradients(zip(grads, [v]))

sok.dump(path,v,optimizer)

sparse_operation_kit.dump_load.load(path, load_vars, optimizer=None)

Abbreviated as sok.load.

Load the embedding tables from sok weight folder. The sok table name must be the same with the weight file prefix.

Now is only support SGD,Adamax,Adadelta,Adagrad,Ftrl,Adam optimizers.

Note: for Adam optimizers , it functions in the same way as LazyAdam in SOK, rather than TensorFlow’s Adam

Parameters:

• path (string) – weight file folder

• load_vars (List,Tuple,SOK Variable) – Can be a single or list of sok.Variable and sok.DynamicVariable

• optimizer (SOK.OptimizerWrapper,optional,default is None) – when model train , need to load optimizer state,input sok.OptimizerWrapper

Return type:

None

Example

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

v = sok.DynamicVariable(dimension=3, var_type="hbm", initializer="13")

optimizer = tf.keras.optimizers.SGD(learning_rate=1.0)
optimizer = sok.OptimizerWrapper(optimizer)
path = "your weight path , weight is dumped by sok.dump"

sok.load(path,v,optimizer)

indices = tf.convert_to_tensor([0, 1, 2**40], dtype=tf.int64)

with tf.GradientTape() as tape:
    embedding = tf.nn.embedding_lookup(v, indices)
    print("embedding:", embedding)
    loss = tf.reduce_sum(embedding)

grads = tape.gradient(loss, [v])
optimizer.apply_gradients(zip(grads, [v]))

sparse_operation_kit.filter_variables(vars)

When using dynamic variables, it is necessary to use sok.OptimizerWrapper to update these variables. Therefore, this API can filter out SOK variables from all TensorFlow variables.

Parameters:

vars (A list of TensorFlow training variables.) –

Returns:

• sok_vars (A list of SOK variables.)

• other_vars (A list of variables that don’t belongs to SOK.)

Example

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

if __name__ == "__main__":
    gpus = tf.config.experimental.list_physical_devices("GPU")
    for gpu in gpus:
        tf.config.experimental.set_memory_growth(gpu, True)

    hvd.init()
    sok.init()

    v1 = tf.Variable([[0, 1, 2]])
    v2 = sok.Variable([[3, 4, 5]])
    v3 = sok.Variable([[6, 7, 8]], mode="localized:0")
    v4 = sok.DynamicVariable(dimension=3, var_type="hbm", initializer="13")

    sok_vars, other_vars = sok.filter_variables([v1, v2, v3, v4])
    assert len(sok_vars) == 3
    assert len(other_vars) == 1

    print("[SOK INFO] filter_variables test passed")