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.