How to do it...

1. To install MXNet on Ubuntu with GPU support, we can use the following command in the terminal:

pip install mxnet-cu80==0.11.0

For other platforms and non-GPU support, have a look at https://mxnet.incubator.apache.org/get_started/install.html.

1. Next, we are ready to import mxnet in our Python environment:

import mxnet as mx

1. We create some simple dummy data that we assign to the GPU and CPU:

import numpy as np
x_input = mx.nd.empty((1, 5), mx.gpu())
x_input[:] = np.array([[1,2,3,4,5]], np.float32)

y_input = mx.nd.empty((1, 5), mx.cpu())
y_input[:] = np.array([[10, 15, 20, 22.5, 25]], np.float32)

1. We can easily copy and adjust the data. Where possible MXNet will automatically execute operations in parallel:

x_input
w_input = x_input
z_input = x_input.copyto(mx.cpu())
x_input += 1
w_input /= 2
z_input *= 2

1. We can print the output as follows:

print(x_input.asnumpy())
print(w_input.asnumpy())
print(z_input.asnumpy())

1. If we want to feed our data to a model, we should create an iterator first:

batch_size = 1
train_iter = mx.io.NDArrayIter(x_input, y_input, batch_size, shuffle=True, data_name='input', label_name='target')

1. Next, we can create the symbols for our model:

X = mx.sym.Variable('input')
Y = mx.symbol.Variable('target')
fc1 = mx.sym.FullyConnected(data=X, name='fc1', num_hidden = 5)
lin_reg = mx.sym.LinearRegressionOutput(data=fc1, label=Y, name="lin_reg")

1. Before we can start training, we need to define our model:

model = mx.mod.Module(
    symbol = lin_reg,
    data_names=['input'], 
    label_names = ['target']
)

1. Let's start training:

model.fit(train_iter,
            optimizer_params={'learning_rate':0.01, 'momentum': 0.9},
            num_epoch=100,
            batch_end_callback = mx.callback.Speedometer(batch_size, 2))

1. To use the trained model for prediction we:

model.predict(train_iter).asnumpy()