MXNet supports Ubuntu AArch64 based operating system so you can run MXNet on all NVIDIA Jetson modules, such as Jetson Nano, TX1, TX2, Xavier NX and AGX Xavier.
These instructions will walk through how to build MXNet and install MXNet's Python language binding.
For the purposes of this install guide we will assume that CUDA is already installed on your Jetson device. NVIDIA Jetpack comes with the latest OS image for Jetson mdoule, and developer tools for both host computer and developer kit, and this also includes CUDA. You should double check what versions are installed and which version you plan to use.
After installing the prerequisites, you have several options for installing MXNet:
To build from source or to use the Python wheel, you must install the following dependencies on your Jetson. Cross-compiling will require dependencies installed on that machine as well.
To use the Python API you need the following dependencies:
sudo apt-get update sudo apt-get install -y \ build-essential \ git \ libopenblas-dev \ libopencv-dev \ python3-pip \ python-numpy sudo pip3 install --upgrade \ pip \ setuptools \ numpy
If you plan to cross-compile you will need to install these dependencies on that computer as well.
These steps are optional, but some of the following instructions expect MXNet source files and the MXNET_HOME
environment variable. Also, CUDA commands will not work out of the box without updating your path.
Clone the MXNet source code repository using the following git
command in your home directory:
git clone --recursive https://github.com/apache/mxnet.git mxnet
You can also checkout a particular branch of MXNet. For example, to install MXNet v1.6:
git clone --recursive -b v1.6.x https://github.com/apache/mxnet.git mxnet
Setup your environment variables for MXNet and CUDA in your .profile
file in your home directory. Add the following to the file.
export PATH=/usr/local/cuda/bin:$PATH export MXNET_HOME=$HOME/mxnet/ export PYTHONPATH=$MXNET_HOME/python:$PYTHONPATH
You can then apply this change immediately with the following:
source .profile
Note: Change the ~/.profile
steps according to how you prefer to use your shell. Otherwise, your environment variables will be gone after you logout.
You can check to see what version of CUDA is running with nvcc
.
nvcc --version
To switch CUDA versions on a device or computer that has more than one version installed, use the following and replace the symbolic link to the version you want. This one uses CUDA 10.2, which comes with Jetpack 4.4.
sudo rm /usr/local/cuda sudo ln -s /usr/local/cuda-10.2 /usr/local/cuda
Note: When cross-compiling, change the CUDA version on the host computer you‘re using to match the version you’re running on your Jetson device.
Installing MXNet from source is a two-step process:
You can use a Docker method or you can build from source manually.
You must have installed Docker and be able to run docker
without sudo
. Follow these setup instructions to get to this point. Then run the following to execute cross-compilation via Docker.
$MXNET_HOME/ci/build.py -p jetson
Step 1 Build the Shared Library
Use the config_jetson.mk Makefile to install MXNet with CUDA bindings to leverage the GPU on the Jetson module.
cp $MXNET_HOME/make/config_jetson.mk config.mk
The pre-existing Makefile builds for all Jetson architectures. Edit config.mk
if you want to specifically build for a particular architecture or if you want to build without CUDA bindings (CPU only). You can make the following changes:
Modify CUDA_ARCH
to build for specific architectures. Currently, we have CUDA_ARCH = -gencode arch=compute_53,code=sm_53 -gencode arch=compute_62,code=sm_62 -gencode arch=compute_72,code=sm_72
. Keep -gencode arch=compute_53,code=sm_53
for Nano and TX1, -gencode arch=compute_62,code=sm_62
for TX2, -gencode arch=compute_72,code=sm_72
for Xavier NX and AGX Xavier.
For CPU only builds, remove USE_CUDA_PATH
, CUDA_ARCH
, USE_CUDNN
flags.
Now you can build the complete MXNet library with the following command:
cd $MXNET_HOME make -j $(nproc)
Executing this command creates a file called libmxnet.so
in the mxnet/lib
directory.
Step 2 Install MXNet Python Bindings (optional)
To install Python bindings run the following commands in the MXNet directory:
cd $MXNET_HOME/python sudo pip3 install -e .
Note that the -e
flag is optional. It is equivalent to --editable
and means that if you edit the source files, these changes will be reflected in the package installed.
Step 3 Install the MXNet Java & Scala Bindings (optional)
Change directories to scala-package
and run mvn install
.
cd $MXNET_HOME/scala-package mvn install
This creates the required .jar
file to use in your Java or Scala projects.
You are now ready to run MXNet on your NVIDIA module. You can verify your MXNet Python installation with the following:
import mxnet mxnet.__version__
You can also verify MXNet can use your GPU with the following test:
import mxnet as mx a = mx.nd.ones((2, 3), mx.gpu()) b = a * 2 + 1 b.asnumpy()
If everything is working, it will report the version number. For assistance, head over to the MXNet Forum.