MXNet: A Scalable Deep Learning Framework

MXNet is an open-source deep learning framework that allows you to define, train, and deploy deep neural networks on a wide array of devices, from cloud infrastructure to mobile devices. It is highly scalable, allowing for fast model training, and supports a flexible programming model and multiple languages. MXNet allows you to mix symbolic and imperative programming flavors to maximize both efficiency and productivity. MXNet is built on a dynamic dependency scheduler that automatically parallelizes both symbolic and imperative operations on the fly. A graph optimization layer on top of that makes symbolic execution fast and memory efficient. The MXNet library is portable and lightweight, and it scales to multiple GPUs and multiple machines.

Please choose the programming language of your choice for the rest of this document.

Installation

Pre-built binaries will be available soon.

Installing the pre-build python package requires a recent version of pip, which, for example, can be installed by

wget https://bootstrap.pypa.io/get-pip.py && sudo python get-pip.py

Will be available soon.

Install by:

pip install mxnet

Use one of following commands to install the desired release:

pip install mxnet       # CPU
pip install mxnet-mkl   # CPU with MKL-DNN acceleration
pip install mxnet-cu75  # GPU with CUDA 7.5
pip install mxnet-cu80  # GPU with CUDA 8.0

The CUDA versions require both CUDA and cuDNN are installed.

Pre-build docker images are available at docker hub.

docker pull mxnet/python
docker pull mxnet/python:gpu
docker pull mxnet/scala
docker pull mxnet/r-lang
docker pull mxnet/r-lang:gpu
docker pull mxnet/julia
docker pull mxnet/julia:gpu

Refer to docker/ for more details.

Refer to the building from source document for details on building MXNet from source codes for various platforms.

Quick Overview

MXNet provides an imperative n-dimensional array interface:

>>> import mxnet as mx
>>> a = mx.nd.ones((2, 3))
>>> b = a * 2 + 1
>>> b.asnumpy()  # print b by converting to a numpy.ndarray object
array([[ 3.,  3.,  3.],
       [ 3.,  3.,  3.]], dtype=float32)
scala> import ml.dmlc.mxnet._
import ml.dmlc.mxnet._

scala> val arr = NDArray.ones(2, 3)
arr: ml.dmlc.mxnet.NDArray = ml.dmlc.mxnet.NDArray@f5e74790

scala> (arr * 2 + 1).toArray
res0: Array[Float] = Array(3.0, 3.0, 3.0, 3.0, 3.0, 3.0)
> require(mxnet)
Loading required package: mxnet
> a <- mx.nd.ones(c(2,3))
> a * 2 + 1
     [,1] [,2] [,3]
[1,]    3    3    3
[2,]    3    3    3
julia> using MXNet

julia> a = mx.ones((2,3))
mx.NDArray{Float32}(2,3)

julia> Array{Float32}(a * 2 + 1)
2×3 Array{Float32,2}:
 3.0  3.0  3.0
 3.0  3.0  3.0
pdl> use AI::MXNet qw(mx)
pdl> $a = mx->nd->ones([2, 3], ctx => mx->gpu())
pdl> print (($a * 2 + 1)->aspdl)
[
 [3 3 3]
 [3 3 3]
]

MXNet also provides a symbolic programming interface:

>>> a = mx.sym.var('a')  # it requires the latest mxnet
>>> b = a * 2 + 1  # b is a Symbol object
>>> c = b.eval(a=mx.nd.ones((2,3)))
>>> c[0].asnumpy()  # the list of outputs
array([[ 3.,  3.,  3.],
       [ 3.,  3.,  3.]], dtype=float32)

Run the above codes in GPU in straightforward:

>>> a = mx.nd.ones((2, 3), mx.gpu())  # create a on GPU 0, then the result a*2+1 will sit on GPU 0 as well
>>> c = b.eval(a=a, ctx=mx.gpu())  # feed a as the input to eval b, the result c will be also on GPU 0
> a <- mx.nd.ones(c(2,3), mx.gpu())
julia> a = mx.ones((2,3), mx.gpu())

In additional, MXNet provides a large number of neural network layers and training modules to facilitate developing deep learning algorithms.

>>> data = mx.sym.var('data')
>>> fc1  = mx.sym.FullyConnected(data, num_hidden=128)
>>> act1 = mx.sym.Activation(fc1, act_type="relu")
>>> fc2  = mx.sym.FullyConnected(act1, num_hidden=10)
>>> loss  = mx.sym.SoftmaxOutput(fc2)
>>> mod = mx.mod.Module(loss)
>>> mod.fit(train_data, ctx=[mx.gpu(0), mx.gpu(1)]) # fit on the training data by using 2 GPUs