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

__init__

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

num_workers

Returns the number of worker nodes.

rank

Returns the rank of this worker node.

type

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