The KV store can accept either string or integer indexes. The mxnet python module uses string keys while interacting with the kv store - for example this line does the init using the param names- python/mxnet/model.py#L100. The optimizer’s interface uses a integer key. The mapping from string name to integer index is maintained in two places -
- within the kv store client code: https://github.com/apache/incubator-mxnet/blob/master/src/kvstore/kvstore_local.h#L82
- computed in the module.init_optimizer https://github.com/apache/incubator-mxnet/blob/master/python/mxnet/module/module.py#L497 and passed onto the optimizer.
In a straightforward use case, both of these mappings will be the same because of the keys are initialized in the same order as in the param_names list. However this can break quite easily - if the kv store has already some keys initialized before being passed to the module.
I think the cleanest solution here is to update the optimizers interface to work with string indexes instead of integers, so that the mapping does not have to be duplicated. Is that something already planned? Is there any other quick fix available? My use case is that i am looking to use the kv store for some auxiliary tasks like computing the mean, std dev of input features of the data across all the workers, adding up the validation loss across the workers etc.