KVStore API

Basic Push and Pull

Basic operation over multiple devices (gpus) on a single machine.

Initialization

Let’s first consider a simple example. It initializes a (int, NDAarray) pair into the store, and then pull the value out.

>>> kv = mx.kv.create('local') # create a local kv store.
>>> shape = (2,3)
>>> kv.init(3, mx.nd.ones(shape)*2)
>>> a = mx.nd.zeros(shape)
>>> kv.pull(3, out = a)
>>> print a.asnumpy()
[[ 2.  2.  2.]
 [ 2.  2.  2.]]

Push, Aggregation, and Updater

For any key has been initialized, we can push a new value with the same shape to the key.

>>> kv.push(3, mx.nd.ones(shape)*8)
>>> kv.pull(3, out = a) # pull out the value
>>> print a.asnumpy()
[[ 8.  8.  8.]
 [ 8.  8.  8.]]

The data for pushing can be on any device. Furthermore, we can push multiple values into the same key, where KVStore will first sum all these values and then push the aggregated value.

>>> gpus = [mx.gpu(i) for i in range(4)]
>>> b = [mx.nd.ones(shape, gpu) for gpu in gpus]
>>> kv.push(3, b)
>>> kv.pull(3, out = a)
>>> print a.asnumpy()
[[ 4.  4.  4.]
 [ 4.  4.  4.]]

For each push, KVStore applies the pushed value into the value stored by a updater. The default updater is ASSGIN, we can replace the default one to control how data is merged.

>>> def update(key, input, stored):
>>>     print "update on key: %d" % key
>>>     stored += input * 2
>>> kv._set_updater(update)
>>> kv.pull(3, out=a)
>>> print a.asnumpy()
[[ 4.  4.  4.]
 [ 4.  4.  4.]]
>>> kv.push(3, mx.nd.ones(shape))
update on key: 3
>>> kv.pull(3, out=a)
>>> print a.asnumpy()
[[ 6.  6.  6.]
 [ 6.  6.  6.]]

Pull

We already see how to pull a single key-value pair. Similar to push, we can also pull the value into several devices by a single call.

>>> b = [mx.nd.ones(shape, gpu) for gpu in gpus]
>>> kv.pull(3, out = b)
>>> print b[1].asnumpy()
[[ 6.  6.  6.]
 [ 6.  6.  6.]]

Interface for list key-value pairs

All operations introduced so far are about a single key. KVStore also provides the interface for a list of key-value pairs. For single device:

>>> keys = [5, 7, 9]
>>> kv.init(keys, [mx.nd.ones(shape)]*len(keys))
>>> kv.push(keys, [mx.nd.ones(shape)]*len(keys))
update on key: 5
update on key: 7
update on key: 9
>>> b = [mx.nd.zeros(shape)]*len(keys)
>>> kv.pull(keys, out = b)
>>> print b[1].asnumpy()
[[ 3.  3.  3.]
 [ 3.  3.  3.]]

For multi-devices:

>>> b = [[mx.nd.ones(shape, gpu) for gpu in gpus]] * len(keys)
>>> kv.push(keys, b)
update on key: 5
update on key: 7
update on key: 9
>>> kv.pull(keys, out = b)
>>> print b[1][1].asnumpy()
[[ 11.  11.  11.]
 [ 11.  11.  11.]]

API Reference

Key value store interface of MXNet for parameter synchronization.

class mxnet.kvstore.KVStore(handle)

A key-value store for synchronization of values, over multiple devices.

init(key, value)

Initialize a single or a sequence of key-value pairs into the store.

For each key, one must init it before push and pull.

Only worker 0’s (rank == 0) data are used.

This function returns after data have been initialized successfully

Parameters:
  • key (int or sequence of int) – The keys.
  • value (NDArray or sequence of NDArray) – The values.

Examples

>>> # init a single key-value pair
>>> shape = (2,3)
>>> kv = mx.kv.create('local')
>>> kv.init(3, mx.nd.ones(shape)*2)
>>> a = mx.nd.zeros(shape)
>>> kv.pull(3, out=a)
>>> print a.asnumpy()
[[ 2.  2.  2.]
[ 2.  2.  2.]]
>>> # init a list of key-value pairs
>>> keys = [5, 7, 9]
>>> kv.init(keys, [mx.nd.ones(shape)]*len(keys))
push(key, value, priority=0)

Push a single or a sequence of key-value pairs into the store.

Data consistency:

  1. this function returns after adding an operator to the engine.

2. push is always called after all previous push and pull on the same key are finished

3. there is no synchronization between workers. One can use _barrier() to sync all workers

Parameters:
  • key (int or list of int) – Keys
  • value (NDArray or list of NDArray or list of list of NDArray) – According values
  • priority (int, optional) – The priority of the push operation. The higher the priority, the faster this action is likely to be executed before other push actions.

Examples

>>> # push a single key-value pair
>>> kv.push(3, mx.nd.ones(shape)*8)
>>> kv.pull(3, out=a) # pull out the value
>>> print a.asnumpy()
[[ 8.  8.  8.]
[ 8.  8.  8.]]
>>> # aggregate the value and the push
>>> gpus = [mx.gpu(i) for i in range(4)]
>>> b = [mx.nd.ones(shape, gpu) for gpu in gpus]
>>> kv.push(3, b)
>>> kv.pull(3, out=a)
>>> print a.asnumpy()
[[ 4.  4.  4.]
[ 4.  4.  4.]]
>>> # push a list of keys.
>>> # single device
>>> kv.push(keys, [mx.nd.ones(shape)]*len(keys))
>>> b = [mx.nd.zeros(shape)]*len(keys)
>>> kv.pull(keys, out=b)
>>> print b[1].asnumpy()
[[ 1.  1.  1.]
[ 1.  1.  1.]]
>>> # multiple devices:
>>> b = [[mx.nd.ones(shape, gpu) for gpu in gpus]] * len(keys)
>>> kv.push(keys, b)
>>> kv.pull(keys, out=b)
>>> print b[1][1].asnumpy()
[[ 4.  4.  4.]
[ 4.  4.  4.]]
pull(key, out=None, priority=0)

Pull a single value or a sequence of values from the store.

Data consistency:

1. this function returns after adding an operator to the engine. But any further read on out will be blocked until it is finished.

2. pull is always called after all previous push and pull on the same key are finished

  1. It pulls the newest value from the store.
Parameters:
  • key (int or list of int) – Keys
  • out (NDArray or list of NDArray or list of list of NDArray) – According values
  • priority (int, optional) – The priority of the push operation. The higher the priority, the faster this action is likely to be executed before other push actions.

Examples

>>> # pull a single key-value pair
>>> a = mx.nd.zeros(shape)
>>> kv.pull(3, out=a)
>>> print a.asnumpy()
[[ 2.  2.  2.]
[ 2.  2.  2.]]
>>> # pull into multiple devices
>>> b = [mx.nd.ones(shape, gpu) for gpu in gpus]
>>> kv.pull(3, out=b)
>>> print b[1].asnumpy()
[[ 2.  2.  2.]
[ 2.  2.  2.]]
>>> # pull a list of key-value pairs.
>>> # On single device
>>> keys = [5, 7, 9]
>>> b = [mx.nd.zeros(shape)]*len(keys)
>>> kv.pull(keys, out=b)
>>> print b[1].asnumpy()
[[ 2.  2.  2.]
[ 2.  2.  2.]]
>>> # On multiple devices
>>> b = [[mx.nd.ones(shape, gpu) for gpu in gpus]] * len(keys)
>>> kv.pull(keys, out=b)
>>> print b[1][1].asnumpy()
[[ 2.  2.  2.]
[ 2.  2.  2.]]
set_optimizer(optimizer)

Register an optimizer to the store

If there are multiple machines, this process (should be a worker node) will pack this optimizer and send it to all servers. It returns after this action is done.

Parameters:optimizer (Optimizer) – the optimizer
type

Get the type of this kvstore

Returns:type – the string type
Return type:str
rank

Get the rank of this worker node

Returns:rank – The rank of this node, which is in [0, get_num_workers())
Return type:int
num_workers

Get the number of worker ndoes

Returns:size – The number of worker nodes
Return type:int
mxnet.kvstore.create(name='local')

Create a new KVStore.

Parameters:name ({'local'}) – The type of KVStore - local works for multiple devices on a single machine (single process) - dist works for multi-machines (multiple processes)
Returns:kv – The created KVStore
Return type:KVStore