Issues with setting MXNet context

[roywei]

I'm creating this issue based on conversation from a keras-mxnet user on changes in latest keras-mxnet:

Setting the context (number of GPUs or CPU) has changed due to the underlying changes in MXNet-1.5.x and this breaks for us. If the context is provided in model.compile(), it is still using the older method, and fails: https://github.com/awslabs/keras-apache-mxnet/blob/master/keras/backend/mxnet_backend.py#L5404 Related to context, there are a few other issues that are problematic:

• For some reason, model.set_mxnet_context() does not allow specifying a single GPU, and I don't see any reason why it should not be allowed. If we're training multiple models on a host with multiple GPUs, this will help us specify different GPUs for different training processes. https://github.com/awslabs/keras-apache-mxnet/blob/master/keras/backend/mxnet_backend.py#L5696
• Another issue existed in the earlier version, as well, but would be good to fix, as we have to work around it currently. The context can be passed in while creating a new model, but "context" is not a supported keyword in the parent class and this results in an exception. https://github.com/awslabs/keras-apache-mxnet/blob/master/keras/backend/mxnet_backend.py#L5348

Here are the reproducible steps.

Using examples/mnist_mlp.py

Test case 1: pass context during compile

Change line45 to the following:

model.compile(loss='categorical_crossentropy',
              optimizer=RMSprop(),
              metrics=['accuracy'], context="gpu(1)")

error

Traceback (most recent call last):
  File "examples/cifar10_resnet.py", line 338, in <module>
    metrics=['accuracy'], context="gpu(0)")
  File "/usr/local/lib/python2.7/dist-packages/keras/backend/mxnet_backend.py", line 5404, in compile
    self._context = self.get_mxnet_context(kwargs['context'])
AttributeError: 'Model' object has no attribute 'get_mxnet_context'

Test case 2: pass context during model initialization

Change line 36 to

model = Sequential(context="gpu(1)")

error

Traceback (most recent call last):
  File "examples/mnist_mlp.py", line 36, in <module>
    model = Sequential(context="gpu(1)")
  File "/usr/local/lib/python2.7/dist-packages/keras/backend/mxnet_backend.py", line 5745, in __init__
    engine.Model.__init__(self, *args, **kwargs)
  File "/usr/local/lib/python2.7/dist-packages/keras/backend/mxnet_backend.py", line 5348, in __init__
    super(Model, self).__init__(*args, **kwargs)
  File "/usr/local/lib/python2.7/dist-packages/keras/legacy/interfaces.py", line 91, in wrapper
    return func(*args, **kwargs)
  File "/usr/local/lib/python2.7/dist-packages/keras/engine/network.py", line 96, in __init__
    self._init_subclassed_network(**kwargs)
TypeError: _init_subclassed_network() got an unexpected keyword argument 'context'

Test case 3: change context using model.set_mxnet_context

add a line before model.compile

model.set_mxnet_context("gpu(1)")

error

Traceback (most recent call last):
  File "examples/mnist_mlp.py", line 44, in <module>
    model.set_mxnet_context("gpu(1)")
  File "/usr/local/lib/python2.7/dist-packages/keras/backend/mxnet_backend.py", line 5710, in set_mxnet_context
    raise ValueError('MXNet Backend: Invalid context provided - %s' % context)
ValueError: MXNet Backend: Invalid context provided - gpu(1)

[roywei]

fixed via https://github.com/awslabs/keras-apache-mxnet/pull/250