#!/bin/bash gpu_sm=7.2 gpu_arch=72 if [ "$1" = "Nano" ] then gpu_sm=5.3 gpu_arch=53 echo "Compile mxnet for sm="$gpu_arch"(Nano)" elif [ "$1" = "TX1" ] then gpu_sm=5.3 gpu_arch=53 echo "Compile mxnet for sm="$gpu_arch"(TX1)" elif [ "$1" = "TX2" ] then gpu_sm=6.2 gpu_arch=62 echo "Compile mxnet for sm="$gpu_arch"(TX2)" elif [ "$1" = "Xavier" ] then gpu_sm=7.2 gpu_arch=72 echo "Compile mxnet for sm="$gpu_arch"(Xavier)" else echo $0" " exit 1 fi # 1. Install dependencies sudo apt-get update sudo apt-get install -y git build-essential python-pip python3-pip ninja-build ccache libprotobuf-dev protobuf-compiler sudo apt-get install -y git libatlas-base-dev libopencv-dev graphviz libopenblas-dev libopenblas-base # 2. Install pip prerequisite sudo pip3 install --upgrade pip sudo pip3 install --upgrade setuptools sudo pip3 install numpy graphviz # 3. Upgrade cmake wget http://www.cmake.org/files/v3.13/cmake-3.13.2.tar.gz tar xpvf cmake-3.13.2.tar.gz cmake-3.13.2/ cd cmake-3.13.2/ ./bootstrap make -j8 sudo make install sudo ldconfig cd .. # 4. Install pybind11 pip3 install pytest git clone https://github.com/pybind/pybind11.git cd pybind11 mkdir build cd build cmake .. -DCMAKE_INSTALL_PREFIX=../install/ make -j8 make install cd ../../ # 5. Apply cublas patch git clone https://github.com/apache/incubator-mxnet.git --branch v1.8.x --recursive mxnet cd mxnet/ git apply ../mxnet_v1.8.x.patch cd .. # 6. Install onnx-onnxtrt export PYVER=3.6 cd mxnet/3rdparty/onnx-tensorrt/third_party/onnx/ export CPLUS_INCLUDE_PATH=/usr/include/python3.6:/usr/local/cuda/targets/aarch64-linux/include mkdir -p build && cd build cmake -DCMAKE_CXX_FLAGS=-I/usr/include/python${PYVER} -DBUILD_ONNX_PYTHON=ON -Dpybind11_DIR=/home/nvidia/pybind11/install/share/cmake/pybind11/ -DBUILD_SHARED_LIBS=ON .. sudo make -j$(nproc) install && \ sudo ldconfig && \ cd .. && \ sudo mkdir -p /usr/include/aarch64-linux-gnu/onnx && \ sudo cp build/onnx/onnx*pb.* /usr/include/aarch64-linux-gnu/onnx && \ sudo cp build/libonnx.so /usr/local/lib && \ sudo rm -f /usr/lib/aarch64-linux-gnu/libonnx_proto.a && \ sudo ldconfig cd ../../ && \ mkdir -p build && \ cd build && \ cmake -DCMAKE_CXX_FLAGS=-I/usr/local/cuda/targets/aarch64-linux/include -DONNX_NAMESPACE=onnx2trt_onnx -DGPU_ARCHS="$gpu_arch" .. && \ sudo make -j$(nproc) install && \ sudo ldconfig cd ../../../../ # 7. Build MXNet cd mxnet mkdir -p build && cd build CUDACXX=/usr/local/cuda/bin/nvcc \ cmake \ -DCMAKE_TOOLCHAIN_FILE=${CMAKE_TOOLCHAIN_FILE} \ -DUSE_CUDA=ON \ -DMXNET_CUDA_ARCH="$gpu_sm" \ -DENABLE_CUDA_RTC=OFF \ -DSUPPORT_F16C=OFF \ -DUSE_OPENCV=ON \ -DUSE_OPENMP=ON \ -DUSE_LAPACK=OFF \ -DUSE_SIGNAL_HANDLER=ON \ -DCMAKE_BUILD_TYPE=Release \ -DUSE_MKL_IF_AVAILABLE=OFF \ -DGPU_ARCHS="$gpu_arch" \ -DUSE_CUDNN=ON \ -DUSE_TENSORRT=ON \ -ONNX_NAMESPACE=onnx \ -DCMAKE_CXX_FLAGS=-I/usr/local/cuda/targets/aarch64-linux/inlude \ .. sudo make -j$(nproc) cd ../.. # 8. Create wheel file set -ex pushd . PYTHON_DIR=$PWD/mxnet/python BUILD_DIR=$PWD/mxnet/build cd ${PYTHON_DIR} python3 setup.py bdist_wheel WHEEL=`readlink -f dist/*.whl` TMPDIR=`mktemp -d` unzip -d ${TMPDIR} ${WHEEL} rm ${WHEEL} cd ${TMPDIR} mv *.data/data/mxnet/libmxnet.so mxnet zip -r ${WHEEL} . cp ${WHEEL} ${BUILD_DIR} rm -rf ${TMPDIR} popd