PyTorch (1.7.1) with CUDA on Drive PX2: Compilation and Installation Guide

Hi,

Following my guides on Compiling Python 3.8.5 on PX2 and Compiling OpenCV 4.5.2 with CUDA, the next important package I needed was PyTorch (with CUDA support). After a lot of trial and errors, a lot of googling and a lot of awesome community answers on how to compile PyTorch with CUDA on PX2 (No cross-compilation => Done directly on PX2), here is a quick guide on how I did it.

Note:

  • PyTorch repo is massive in size and the files generated during the build process further exaggerate the space occupied. Therefore, it is recommended to have atleast 3 GB+ free space.
  • Many requied deps knowingly or unknowingly maybe covered in my above mentioned guides of Python and OpenCV. So it is better to execute these steps and follow up here.
  • During the build process, I have noticed number of hang ups and crashes caused by lack of memory (not storage). To prevent this issue, create a swap file for atleast 2 GB (I did 4 GB) or use less number of cores (<4) while building

Installation:

  • Install multi-processing and image-processing libraries
sudo apt-get install libblas-dev liblapack-dev libopenblas-dev libopenmpi-dev libomp-dev libjpeg-dev zlib1g-dev libpython3-dev libavcodec-dev libavformat-dev libswscale-dev
  • Compile and install Cmake (3.22)
wget https://cmake.org/files/v3.22/cmake-3.22.5.tar.gz
cd cmake-3.22.5/
./configure
make
sudo make install
  • (Optional but recommended) Create a python virtual enviroment (or if already existing, replace the paths in build commands)
python3.8 -m venv venv
source venv/bin/activate
  • Update pip and install few key python packages
python -m pip install -U pip
python -m pip install numpy matplotlib opencv-python ninja scikit-build pyyaml
  • Clone PyTorch and switch branch to v1.7.1 (latest PyTorch version supporting CUDA 9.2)
git clone --recursive --branch v1.7.1 http://github.com/pytorch/pytorch
cd pytorch
pip3 install -r requirements.txt
  • Compile and install PyTorch (Entire build process can take 20+ hours on PX2🥴)
python setup.py clean
USE_FBGEMM=0 NO_DISTRIBUTED=1 USE_DISTRIBUTED=0 NO_MKLDNN=1 NO_XNNPACK=1 USE_QNNPACK=0 USE_PYTORCH_QNNPACK=0 MAX_JOBS=3 TORCH_CUDA_ARCH_LIST='5.3;6.1;6.2' BUILD_BINARY=1 USE_NCCL=0 BUILD_TEST=0 python setup.py install
  • (Optional, not sure about stability) Compile and install TorchVision (0.8.2)
git clone --recursive --branch v0.8.2 https://github.com/pytorch/vision.git
cd vision
python setup.py clean
USE_FBGEMM=0 NO_DISTRIBUTED=1 USE_DISTRIBUTED=0 NO_MKLDNN=1 NO_XNNPACK=1 USE_QNNPACK=0 USE_PYTORCH_QNNPACK=0 MAX_JOBS=3 TORCH_CUDA_ARCH_LIST='5.3;6.1;6.2' BUILD_BINARY=1 USE_NCCL=0 BUILD_TEST=0 python setup.py install

Finally done🥳🥳 Now you can use PyTorch (and possibly also TorchVision) in Python or as an added advantage LibTorch using the C++ API with full CUDA support. A small parting-tip from fellow Drive PX2 AI developer: don’t use FP16, PX2 doesn’t support it and it can potentially slow down your inference

These commands were tested and worked as of Mai 2022. Hope it helps !!
Any suggestions/improvements are highly appreciated.

Cheers,
Vishal Balaji
Perception Engineer@Schanzer Racing

Dear @vishal.balaji,
Thank you for sharing PyTorch installation steps to help others in community.