NDArray - Scientific computing on CPU and GPU¶
NDArray is a tensor data structure similar to numpy’s multi-dimensional array. In addition, it supports asynchronous computation on CPU and GPU.
First, let’s import MXNet:
from __future__ import print_function import numpy as np import mxnet as mx
Creating NDArray¶
There are many ways to create NDArray.
Construct from (nested) list:
x = mx.nd.array([[1, 2, 3], [4, 5, 6]]) print(x)
Construct from numpy array:
x_numpy = np.ones((2, 3)) x = mx.nd.array(x_numpy) print(x)
Array construction routines:
# create an 2x3 array of ones x = mx.nd.ones((2, 3)) print(x) # create an 2x3 array of zeros x = mx.nd.zeros((2, 3)) print(x) # create an 1d-array of 0 to 5 and reshape to 2x3 x = mx.nd.arange(6).reshape((2, 3)) print(x)
You can convert an NDArray to numpy array to retrieve its data with
.asnumpy():
z = x.asnumpy() print(z)
Basic attributes¶
NDArray has some basic attributes that you often want to query:
NDArray.shape: The dimensions of the array. It is a tuple of integers
indicating the length of the array along each axis. For a matrix with
n rows
and
m columns, its
shape will be
(n, m).
print('x.shape:', x.shape)
NDArray.dtype: A
numpy type object describing the type of array
elements.
print('x.dtype:', x.dtype)
NDArray.size: the total number of components in the array - equals to the
product of the components of its
shape
print('x.size:', x.size)
NDArray.context: The device on which this array is stored, e.g.
mx.cpu()
or
mx.gpu(1).
print('x.context:', x.context)
NDArray Operations¶
NDArray supports a wide range of operations. Simple operations can be called with python syntax:
x = mx.nd.array([[1, 2], [3, 4]]) y = mx.nd.array([[4, 3], [2, 1]]) print(x + y)
You can also call operators from the
mxnet.ndarray (or
mx.nd for short) name space:
z = mx.nd.add(x, y) print(z)
You can also pass additional flags to operators:
z = mx.nd.sum(x, axis=0) print('axis=0:', z) z = mx.nd.sum(x, axis=1) print('axis=1:', z)
Using GPU¶
Each NDArray lives on a
Context. MXNet supports
mx.cpu() for CPU and
mx.gpu(0),
mx.gpu(1), etc for GPU. You can specify context when creating NDArray:
# creates on CPU (the default). # Replace mx.cpu() with mx.gpu(0) if you have a GPU. x = mx.nd.zeros((2, 2), ctx=mx.cpu()) print(x)
x = mx.nd.array([[1, 2], [3, 4]], ctx=mx.cpu()) print(x)
You can copy arrays between devices with
.copyto():
# Copy x to cpu. Replace with mx.gpu(0) if you have GPU. y = x.copyto(mx.cpu()) print(y)
# Copy x to another NDArray, possibly on another Context. y = mx.nd.zeros_like(x) x.copyto(y) print(y)
See the Advanced NDArray tutorial for a more detailed introduction to NDArray API.