mxnet.kvstore.byteps

BytePS backend for MXNet KVStore

Classes

BytePS()

BytePS backend for MXNet KVStore interface.

class mxnet.kvstore.byteps.BytePS[source]

Bases: KVStoreBase

BytePS backend for MXNet KVStore interface.

broadcast(key, value, out, priority=0)[source]

Broadcast the value NDArray at rank 0 to all ranks’ out. If out is None, the result is stored in value.

Parameters:

Examples

>>> # broadcast a single key-value pair
>>> shape = (2,3)
>>> kv = mx.kv.create('byteps')
>>> a = mx.nd.zeros(shape)
>>> kv.broadcast('3', mx.nd.ones(shape)*2, out=a)
>>> print a.asnumpy()
[[ 2.  2.  2.]
[ 2.  2.  2.]]
static is_capable(capability)[source]

Queries if the KVStore type supports certain capability, such as optimizer algorithm, gradient compression, sparsity, etc. As byteps server does not store weight, this function will return false for any capabilities.

Parameters:

capability (str) – The capability to query

Returns:

result – Whether the capability is supported or not.

Return type:

bool

load_optimizer_states(fname)[source]

Not Implement yet.

Parameters:

fname (str) – Path to input states file.

property local_rank

Returns the local rank of this worker on the node.

Returns:

rank – The local rank of this node, which is in range [0, num_workers_on_current_node())

Return type:

int

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 from the store. This function is coalesced form of push and pull operations. value is pushed to the kvstore server for the specified keys and the aggregated values are pulled from the server to out. If out is not specified the pulled values are written to value.

Parameters:
  • key (str, or int) – The key.

  • value (NDArray, or list of NDArray) – Values corresponding to the key.

  • 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.

Examples

>>> # pushpull a single key-value pair
>>> kv.pushpull('3', mx.nd.ones(shape)*8, out=a)
>>> print a.asnumpy()
[[ 8.  8.  8.]
[ 8.  8.  8.]]
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

save_optimizer_states(fname, dump_optimizer=False)[source]

Not Implement yet.

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]

Not Implement yet.

Parameters:

optimizer (KVStoreBase) – The new optimizer for the store

property type

Returns the type of this kvstore.

Returns:

type – the string type

Return type:

str