SparseOperationKit Initialize

sparse_operation_kit.init(comm_tool='horovod', use_legacy_optimizer=True)[source]

Abbreviated as sok.init.

This function is used to do the initialization of SparseOperationKit (SOK).

SOK will leverage all available GPUs for current CPU process. Please set CUDA_VISIBLE_DEVICES or tf.config.set_visible_devices to specify which GPU(s) are used in this process before launching tensorflow runtime and calling this function.

Currently, these API only support horovod as the communication tool, so horovod.init must be called before initializing SOK.

Parameters:

  • comm_tool (string) – a string to specify which communication tool to use. Default value is “horovod”.

  • use_legacy_optimizer (bool) – From tensorflow 2.11.0 , keras default optimizer is optimizer experimental. SOK won’t support it in future, so if you switch use_legacy_optimizer to True, SOK will redefine tensorflow.keras.optimizers to tensorflow.keras.optimizers.legacy(tf.keras.optimizers.optimizer_v2). Default value is True, if you want to use new optimizer in the other part in your code , and only use legacy optimizer in SOK, please set to False

Return type:

None

Example

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()