Manipulate data the MXNet way with ndarray

It’s impossible to get anything done if we can’t manipulate data. This has two parts - loading data and processing data once it’s inside the computer. This notebook is about the latter. So let’s start by introducing NDArrays, MXNet’s primary tool for storing and transforming data. If you’ve worked with NumPy before, you’ll notice that NDArrays are by design similar to NumPy’s multi-dimensional array. However, they confer a few key advantages. First, NDArrays support asynchronous computation on CPU, GPU, and distributed cloud architectures. Second, they provide support for automatic differentiation. These properties make NDArray an ideal library for machine learning, both for researchers and engineers launching production systems.

Getting started

In this chapter, we’ll get you going with the basic functionality. Don’t worry if you don’t understand any of the basic math, like element-wise operations or normal distributions. In the next two chapters we’ll take another pass at NDArray, teaching you both the math you’ll need and how to realize it in code.

To get started, let’s import mxnet. We’ll also import ndarray from mxnet for convenience. We’ll make a habit of setting a random seed so that you always get the same results that we do.

In [1]:
import mxnet as mx
from mxnet import nd
mx.random.seed(1)

Next, let’s see how to create an NDArray, without any values initialized. Specifically, we’ll create a 2D array (also called a matrix) with 3 rows and 4 columns.

In [2]:
x = nd.empty((3, 4))
print(x)

[[  0.00000000e+00   0.00000000e+00  -3.38584985e+16   4.58462819e-41]
 [  1.38654559e-38   0.00000000e+00  -1.25588698e+32  -4.44981180e+24]
 [ -8.66752053e+18   4.58462819e-41  -9.26639042e+22   4.58448806e-41]]
<NDArray 3x4 @cpu(0)>

The empty method just grabs some memory and hands us back a matrix without setting the values of any of its entries. This means that the entries can have any form of values, including very big ones! But typically, we’ll want our matrices initialized. Commonly, we want a matrix of all zeros.

In [3]:
x = nd.zeros((3, 5))
x
Out[3]:

[[ 0.  0.  0.  0.  0.]
 [ 0.  0.  0.  0.  0.]
 [ 0.  0.  0.  0.  0.]]
<NDArray 3x5 @cpu(0)>

Similarly, ndarray has a function to create a matrix of all ones.

In [4]:
x = nd.ones((3, 4))
x
Out[4]:

[[ 1.  1.  1.  1.]
 [ 1.  1.  1.  1.]
 [ 1.  1.  1.  1.]]
<NDArray 3x4 @cpu(0)>

Often, we’ll want to create arrays whose values are sampled randomly. This is especially common when we intend to use the array as a parameter in a neural network. In this snippet, we initialize with values drawn from a standard normal distribution with zero mean and unit variance.

In [5]:
y = nd.random_normal(0, 1, shape=(3, 4))
y
Out[5]:

[[-0.67765152  0.03629481  0.10073948 -0.49024421]
 [ 0.57595438 -0.95017916 -0.3469252   0.03751944]
 [-0.22134334 -0.72984636 -1.80471897 -2.04010558]]
<NDArray 3x4 @cpu(0)>

As in NumPy, the dimensions of each NDArray are accessible via the .shape attribute.

In [6]:
y.shape
Out[6]:
(3, 4)

We can also query its size, which is equal to the product of the components of the shape. Together with the precision of the stored values, this tells us how much memory the array occupies.

In [7]:
y.size
Out[7]:
12

Operations

NDArray supports a large number of standard mathematical operations. Such as element-wise addition:

In [8]:
x + y
Out[8]:

[[ 0.32234848  1.03629482  1.10073948  0.50975579]
 [ 1.57595444  0.04982084  0.6530748   1.03751945]
 [ 0.77865666  0.27015364 -0.80471897 -1.04010558]]
<NDArray 3x4 @cpu(0)>

Multiplication:

In [9]:
x * y
Out[9]:

[[-0.67765152  0.03629481  0.10073948 -0.49024421]
 [ 0.57595438 -0.95017916 -0.3469252   0.03751944]
 [-0.22134334 -0.72984636 -1.80471897 -2.04010558]]
<NDArray 3x4 @cpu(0)>

And exponentiation:

In [10]:
nd.exp(y)
Out[10]:

[[ 0.50780815  1.03696156  1.1059885   0.61247683]
 [ 1.77882743  0.38667175  0.70685822  1.03823221]
 [ 0.80144149  0.48198304  0.16452068  0.13001499]]
<NDArray 3x4 @cpu(0)>

We can also grab a matrix’s transpose to compute a proper matrix-matrix product.

In [11]:
nd.dot(x, y.T)
Out[11]:

[[-1.03086138 -0.68363053 -4.79601431]
 [-1.03086138 -0.68363053 -4.79601431]
 [-1.03086138 -0.68363053 -4.79601431]]
<NDArray 3x3 @cpu(0)>

We’ll explain these operations and present even more operators in the linear algebra chapter. But for now, we’ll stick with the mechanics of working with NDArrays.

In-place operations

In the previous example, every time we ran an operation, we allocated new memory to host its results. For example, if we write y = x + y, we will dereference the matrix that y used to point to and instead point it at the newly allocated memory. We can show this using Python’s id() function, which tells us precisely which object a variable refers to.

In [12]:
print('id(y):', id(y))
y = y + x
print('id(y):', id(y))
id(y): 140522018446248
id(y): 140518002409312

We can assign the result to a previously allocated array with slice notation, e.g., result[:] = ....

In [13]:
z = nd.zeros_like(x)
print('id(z):', id(z))
z[:] = x + y
print('id(z):', id(z))
id(z): 140522018501240
id(z): 140522018501240

However, x+y here will still allocate a temporary buffer to store the result before copying it to z. To make better use of memory, we can perform operations in place, avoiding temporary buffers. To do this we specify the out keyword argument every operator supports:

In [14]:
nd.elemwise_add(x, y, out=z)
Out[14]:

[[ 1.32234848  2.03629494  2.10073948  1.50975585]
 [ 2.57595444  1.0498209   1.65307474  2.03751945]
 [ 1.77865672  1.27015364  0.19528103 -0.04010558]]
<NDArray 3x4 @cpu(0)>

If we’re not planning to re-use x, then we can assign the result to x itself. There are two ways to do this in MXNet. 1. By using slice notation x[:] = x op y 2. By using the op-equals operators like +=

In [15]:
print('id(x):', id(x))
x += y
x
print('id(x):', id(x))
id(x): 140522018503144
id(x): 140522018503144

Slicing

MXNet NDArrays support slicing in all the ridiculous ways you might imagine accessing your data. Here’s an example of reading the second and third rows from x.

In [16]:
x[1:2]
Out[16]:

[[ 2.57595444  1.0498209   1.65307474  2.03751945]]
<NDArray 1x4 @cpu(0)>

Now let’s try writing to a specific element.

In [17]:
x[1,2] = 9.0
x
Out[17]:

[[ 1.32234848  2.03629494  2.10073948  1.50975585]
 [ 2.57595444  1.0498209   9.          2.03751945]
 [ 1.77865672  1.27015364  0.19528103 -0.04010558]]
<NDArray 3x4 @cpu(0)>

Multi-dimensional slicing is also supported.

In [18]:
x[1:2,1:3]
Out[18]:

[[ 1.0498209  9.       ]]
<NDArray 1x2 @cpu(0)>
In [19]:
x[1:2,1:3] = 5.0
x
Out[19]:

[[ 1.32234848  2.03629494  2.10073948  1.50975585]
 [ 2.57595444  5.          5.          2.03751945]
 [ 1.77865672  1.27015364  0.19528103 -0.04010558]]
<NDArray 3x4 @cpu(0)>

Broadcasting

You might wonder, what happens if you add a vector y to a matrix X? These operations, where we compose a low dimensional array y with a high-dimensional array X invoke a functionality called broadcasting. Here, the low-dimensional array is duplicated along any axis with dimension \(1\) to match the shape of the high dimensional array. Consider the following example.

In [20]:
x = nd.ones(shape=(3,3))
print('x = ', x)
y = nd.arange(3)
print('y = ', y)
print('x + y = ', x + y)
x =
[[ 1.  1.  1.]
 [ 1.  1.  1.]
 [ 1.  1.  1.]]
<NDArray 3x3 @cpu(0)>
y =
[ 0.  1.  2.]
<NDArray 3 @cpu(0)>
x + y =
[[ 1.  2.  3.]
 [ 1.  2.  3.]
 [ 1.  2.  3.]]
<NDArray 3x3 @cpu(0)>

While y is initially of shape (3), MXNet infers its shape to be (1,3), and then broadcasts along the rows to form a (3,3) matrix). You might wonder, why did MXNet choose to interpret y as a (1,3) matrix and not (3,1). That’s because broadcasting prefers to duplicate along the left most axis. We can alter this behavior by explicitly giving y a 2D shape.

In [21]:
y = y.reshape((3,1))
print('y = ', y)
print('x + y = ', x+y)
y =
[[ 0.]
 [ 1.]
 [ 2.]]
<NDArray 3x1 @cpu(0)>
x + y =
[[ 1.  1.  1.]
 [ 2.  2.  2.]
 [ 3.  3.  3.]]
<NDArray 3x3 @cpu(0)>

Converting from MXNet NDArray to NumPy

Converting MXNet NDArrays to and from NumPy is easy. The converted arrays do not share memory.

In [22]:
a = x.asnumpy()
type(a)
Out[22]:
numpy.ndarray
In [23]:
y = nd.array(a)
y
Out[23]:

[[ 1.  1.  1.]
 [ 1.  1.  1.]
 [ 1.  1.  1.]]
<NDArray 3x3 @cpu(0)>

Managing context

You might have noticed that MXNet NDArray looks almost identical to NumPy. But there are a few crucial differences. One of the key features that differentiates MXNet from NumPy is its support for diverse hardware devices.

In MXNet, every array has a context. One context could be the CPU. Other contexts might be various GPUs. Things can get even hairier when we deploy jobs across multiple servers. By assigning arrays to contexts intelligently, we can minimize the time spent transferring data between devices. For example, when training neural networks on a server with a GPU, we typically prefer for the model’s parameters to live on the GPU. To start, let’s try initializing an array on the first GPU.

In [24]:
from mxnet import gpu
z = nd.ones(shape=(3,3), ctx=gpu(0))
z
Out[24]:

[[ 1.  1.  1.]
 [ 1.  1.  1.]
 [ 1.  1.  1.]]
<NDArray 3x3 @gpu(0)>

Given an NDArray on a given context, we can copy it to another context by using the copyto() method.

In [25]:
x_gpu = x.copyto(gpu(0))
print(x_gpu)

[[ 1.  1.  1.]
 [ 1.  1.  1.]
 [ 1.  1.  1.]]
<NDArray 3x3 @gpu(0)>

The result of an operator will have the same context as the inputs.

In [26]:
x_gpu + z
Out[26]:

[[ 2.  2.  2.]
 [ 2.  2.  2.]
 [ 2.  2.  2.]]
<NDArray 3x3 @gpu(0)>

Watch out!

Imagine that your variable z already lives on your second GPU (gpu(0)). What happens if we call z.copyto(gpu(0))? It will make a copy and allocate new memory, even though that variable already lives on the desired device!

Often, we only want to make a copy if the variable currently lives in the wrong context. In these cases, we can call as_in_context(). If the variable is already on gpu(0) then this is a no-op.

In [29]:
print('id(z):', id(z))
z = z.copyto(gpu(0))
print('id(z):', id(z))
z = z.as_in_context(gpu(0))
print('id(z):', id(z))
print(z)
id(z): 140518002062728
id(z): 140518002063456
id(z): 140518002063456

[[ 1.  1.  1.]
 [ 1.  1.  1.]
 [ 1.  1.  1.]]
<NDArray 3x3 @gpu(0)>