SparseOperationKit Initialize

sparse_operation_kit.core.initialize.Init(**kwargs)[source]

Abbreviated as sok.Init(**kwargs).

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.

In TensorFlow 2.x, SOK can be used with tf.distribute.Strategy or Horovod. When it’s used with tf.distribute.Strategy, it must be called under strategy.scope(). For example,

with strategy.scope():
    sok.Init(**kwargs)

When it’s used with Horovod, it must be called at each process. For example,

import horovod.tensorflow as hvd

hvd.init()

sok.Init(**kwargs)

In TensorFlow 1.15, SOK can only work with Horovod. The retured status must be evaluated with sess.run, and it must be the first step before evaluate any other SOK APIs.

sok_init = sok.Init(global_batch_size=args.global_batch_size)
with tf.Session() as sess:
    sess.run(sok_init)
    ...

Parameters: kwargs (dictionary) – keyword arguments for this function. Currently, it must contains global_batch_size used in all GPUs.
Returns: status – a string will be returned if this function executed successfully. And its contents will be ‘OK’.
Return type: string