mxnet.kvstore.KVStoreBase¶
-
class
KVStoreBase
[source]¶ Bases:
object
An abstract key-value store interface for data parallel training.
-
__init__
()¶ Initialize self. See help(type(self)) for accurate signature.
Methods
Initialize self.
broadcast
(key, value, out[, priority])Broadcast the value NDArray at rank 0 to all ranks, and store the result in out
is_capable
(capability)Queries if the KVStore type supports certain capability, such as optimizer algorithm, gradient compression, sparsity, etc.
load_optimizer_states
(fname)Loads the optimizer (updater) state from the file.
pushpull
(key, value[, out, priority])Performs push and pull a single value or a sequence of values from the store.
register
(klass)Registers a new KVStore.
save_optimizer_states
(fname[, dump_optimizer])Saves the optimizer (updater) state to a file.
set_optimizer
(optimizer)Registers an optimizer with the kvstore.
Attributes
OPTIMIZER
kv_registry
Returns the number of worker nodes.
Returns the rank of this worker node.
Returns the type of this kvstore backend.
-
broadcast
(key, value, out, priority=0)[source]¶ Broadcast the value NDArray at rank 0 to all ranks, and store the result in out
- Parameters
key (str or int) – The key.
value (NDArray) – The value corresponding to the key to broadcast
out (NDArray, or list of NDArray) – Values corresponding to the key to store the result
priority (int, optional) – The priority of the operation. Higher priority operations are likely to be executed before other actions.
-
is_capable
(capability)[source]¶ Queries if the KVStore type supports certain capability, such as optimizer algorithm, gradient compression, sparsity, etc.
- Parameters
capability (str) – The capability to query
- Returns
result – Whether the capability is supported or not.
- Return type
bool
-
load_optimizer_states
(fname)[source]¶ Loads the optimizer (updater) state from the file.
- Parameters
fname (str) – Path to input states file.
-
property
num_workers
¶ Returns the number of worker nodes.
- Returns
size – The number of worker nodes.
- Return type
int
-
pushpull
(key, value, out=None, priority=0)[source]¶ Performs push and pull a single value or a sequence of values from the store.
This function is coalesced form of push and pull operations.
value is pushed to the kvstore server for summation with the specified keys, and the results are pulled from the server to out. If out is not specified the pulled values are written to value.
Note that for allreduce based approaches such as horovod, there is no notion of server or store. This function performs allreduce.
- Parameters
key (str or int) – The key.
value (NDArray, or list of NDArray) – Values corresponding to the keys.
out (NDArray, or list of NDArray) – Values corresponding to the key.
priority (int, optional) – The priority of the operation. Higher priority operations are likely to be executed before other actions.
-
property
rank
¶ Returns the rank of this worker node.
- Returns
rank – The rank of this node, which is in range [0, num_workers())
- Return type
int
-
static
register
(klass)[source]¶ Registers a new KVStore. Once a kvstore is registered, we can create an instance of this kvstore with create later.
Examples
>>> @mx.kvstore.KVStoreBase.register ... class MyKVStore(mx.kvstore.KVStoreBase): ... pass >>> kv = mx.kv.create('MyKVStore') >>> print(type(kv)) <class '__main__.MyKVStore'>
-
save_optimizer_states
(fname, dump_optimizer=False)[source]¶ Saves the optimizer (updater) state to a file. This is often used when checkpointing the model during training.
- Parameters
fname (str) – Path to the output states file.
dump_optimizer (bool, default False) – Whether to also save the optimizer itself. This would also save optimizer information such as learning rate and weight decay schedules.
-
set_optimizer
(optimizer)[source]¶ Registers an optimizer with the kvstore.
When using a single machine, this function updates the local optimizer. If using multiple machines and this operation is invoked from a worker node, it will serialized the optimizer with pickle and send it to all servers. The function returns after all servers have been updated.
- Parameters
optimizer (KVStoreBase) – The new optimizer for the store
-
property
type
¶ Returns the type of this kvstore backend.
- Returns
type – the string type
- Return type
str
-