NDArray API

The NDArray package (mxnet.ndarray) contains tensor operations similar to numpy.ndarray. The syntax is also similar, except for some additional calls for dealing with I/O and multiple devices.

Topics:

Create NDArray

Create mxnet.ndarray as follows:

import org.apache.mxnet._
// all-zero array of dimension 100x50
val a = NDArray.zeros(100, 50)
// all-one array of dimension 256x32x128x1
val b = NDArray.ones(256, 32, 128, 1)
// initialize array with contents, you can specify dimensions of array using Shape parameter while creating array.
val c = NDArray.array(Array(1, 2, 3, 4, 5, 6), shape = Shape(2, 3))

This is similar to the way you use numpy.

NDArray Operations

We provide some basic ndarray operations, like arithmetic and slice operations.

Arithmetic Operations

import org.apache.mxnet._
val a = NDArray.zeros(100, 50)
a.shape
// org.apache.mxnet.Shape = (100,50)
val b = NDArray.ones(100, 50)
// c and d will be calculated in parallel here!
val c = a + b
val d = a - b
// inplace operation, b's contents will be modified, but c and d won't be affected.
b += d

Multiplication/Division Operations

import org.apache.mxnet._
// Multiplication
val ndones = NDArray.ones(2, 1)
val ndtwos = ndones * 2
ndtwos.toArray
// Array[Float] = Array(2.0, 2.0)
(ndones * ndones).toArray
// Array[Float] = Array(1.0, 1.0)
(ndtwos * ndtwos).toArray
// Array[Float] = Array(4.0, 4.0)
ndtwos *= ndtwos // inplace
ndtwos.toArray
// Array[Float] = Array(4.0, 4.0)

//Division
val ndones = NDArray.ones(2, 1)
val ndzeros = ndones - 1f
val ndhalves = ndones / 2
ndhalves.toArray
// Array[Float] = Array(0.5, 0.5)
(ndhalves / ndhalves).toArray
// Array[Float] = Array(1.0, 1.0)
(ndones / ndones).toArray
// Array[Float] = Array(1.0, 1.0)
(ndzeros / ndones).toArray
// Array[Float] = Array(0.0, 0.0)
ndhalves /= ndhalves
ndhalves.toArray
// Array[Float] = Array(1.0, 1.0)

Slice Operations

import org.apache.mxnet._
val a = NDArray.array(Array(1f, 2f, 3f, 4f, 5f, 6f), shape = Shape(3, 2))
val a1 = a.slice(1)
assert(a1.shape === Shape(1, 2))
assert(a1.toArray === Array(3f, 4f))

val a2 = arr.slice(1, 3)
assert(a2.shape === Shape(2, 2))
assert(a2.toArray === Array(3f, 4f, 5f, 6f))

Dot Product

import org.apache.mxnet._
val arr1 = NDArray.array(Array(1f, 2f), shape = Shape(1, 2))
val arr2 = NDArray.array(Array(3f, 4f), shape = Shape(2, 1))
val res = NDArray.dot(arr1, arr2)
res.shape
// org.apache.mxnet.Shape = (1,1)
res.toArray
// Array[Float] = Array(11.0)

Save and Load NDArray

You can use MXNet functions to save and load a list or dictionary of NDArrays from file systems, as follows:

import org.apache.mxnet._
val a = NDArray.zeros(100, 200)
val b = NDArray.zeros(100, 200)
// save list of NDArrays
NDArray.save("/path/to/array/file", Array(a, b))
// save dictionary of NDArrays to AWS S3
NDArray.save("s3://path/to/s3/array", Map("A" -> a, "B" -> b))
// save list of NDArrays to hdfs.
NDArray.save("hdfs://path/to/hdfs/array", Array(a, b))
val from_file = NDArray.load("/path/to/array/file")
val from_s3 = NDArray.load("s3://path/to/s3/array")
val from_hdfs = NDArray.load("hdfs://path/to/hdfs/array")

The good thing about using the save and load interface is that you can use the format across all mxnet language bindings. They also already support Amazon S3 and HDFS.

Multi-Device Support

Device information is stored in the mxnet.Context structure. When creating NDArray in MXNet, you can use the context argument (the default is the CPU context) to create arrays on specific devices as follows:

import org.apache.mxnet._
val cpu_a = NDArray.zeros(100, 200)
cpu_a.context
// org.apache.mxnet.Context = cpu(0)
val ctx = Context.gpu(0)
val gpu_b = NDArray.zeros(Shape(100, 200), ctx)
gpu_b.context
// org.apache.mxnet.Context = gpu(0)

Currently, we do not allow operations among arrays from different contexts. To manually enable this, use the copyto member function to copy the content to different devices, and continue computation:

import org.apache.mxnet._
val x = NDArray.zeros(100, 200)
val ctx = Context.gpu(0)
val y = NDArray.zeros(Shape(100, 200), ctx)
val z = x + y
// mxnet.base.MXNetError: [13:29:12] src/ndarray/ndarray.cc:33:
// Check failed: lhs.ctx() == rhs.ctx() operands context mismatch
val cpu_y = NDArray.zeros(100, 200)
y.copyto(cpu_y)
val z = x + cpu_y

Next Steps

  • See KVStore API for multi-GPU and multi-host distributed training.