Skip to content

Instantly share code, notes, and snippets.

@mmattamala
Last active July 4, 2022 16:08
Show Gist options
  • Save mmattamala/e1725bb521e5d794eb5b107bd09a5997 to your computer and use it in GitHub Desktop.
Save mmattamala/e1725bb521e5d794eb5b107bd09a5997 to your computer and use it in GitHub Desktop.
Pytorch 1.10 compilation on Xavier

General instructions

https://forums.developer.nvidia.com/t/pytorch-for-jetson-version-1-11-now-available/72048

Increase performance

sudo nvpmodel -m 2     # on Xavier NX, use -m 2  instead (15W 6-core mode)
sudo jetson_clocks

Download pytorch

git clone --recursive --branch release/1.10 http://github.com/pytorch/pytorch
cd pytorch

Apply patch for 1.10

wget https://gist.githubusercontent.com/dusty-nv/ce51796085178e1f38e3c6a1663a93a1/raw/4f1a0f948150c91f877aa38075835df748c81fe5/pytorch-1.10-jetpack-4.5.1.patch
git apply pytorch-1.10-jetpack-4.5.1.patch

Export env variables

export USE_NCCL=0
export USE_DISTRIBUTED=0 # skip setting this if you want to enable OpenMPI backend
export USE_QNNPACK=0
export USE_PYTORCH_QNNPACK=0
export TORCH_CUDA_ARCH_LIST="7.2" # for Xavier

export PYTORCH_BUILD_VERSION=1.10.0  # without the leading 'v', e.g. 1.3.0 for PyTorch v1.3.0
export PYTORCH_BUILD_NUMBER=1

Install dependencies

sudo apt-get install python3-pip cmake libopenblas-dev libopenmpi-dev 

pip3 install -r requirements.txt
pip3 install scikit-build
pip3 install ninja

Build

python3 setup.py bdist_wheel

Note: Downgrade setuptools

If you get AttributeError: module 'setuptools._distutils' has no attribute 'version'., downgrade setuptools (reference):

pip install setuptools==59.5.0
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment