Jetson install GPU PyTorch

Synopsis

To whom needs GPU version of PyTorch running on NVIDIA Jetson.

Prerequisite checklist

1. Check model of embedded board
2. Prepare a virtual python environment (e.g., miniforge3), it’s not suggested to follow NVIDIA official version of guideline to install code by sudo (which may mess up system python environments).
3. Check pre-built binary files on Jetson binaries, if there are some appropriate prebuilt wheels, then download them. Make sure to rename them into the formate like “ torch-1.10.0-cp36-cp36m-linux_aarch64.whl”.
4. Run the following code:

   sudo apt-get update
   sudo apt install python3-dev python3-pip
   sudo apt-get install libhdf5-serial-dev hdf5-tools libhdf5-dev zlib1g-dev zip libjpeg8-dev liblapack-dev libblas-dev gfortran
   sudo apt-get install python3-pip
   # for Jetpack version >= 5.0.2, python 3.8.10 is suggested.
   conda create -n pytorch_env python=3.8.10
   conda activate pytorch_env
   pip install pip testresources setuptools
   # install by wheel file (example for newest Jetpack 5.1.2)
   # News: PyTorch 2.0.0 prebuilt is supported right now.
   pip install torch-2.0.0+nv23.05-cp38-cp38-linux_aarch64.whl

Install by source code (not suggested)

Unless you really need some typical version of PyTorch, otherwise to build a PyTorch wheel is really unnecessary (time-consuming, bug-filled, lack of documents)

sudo apt-get update
sudo apt-get upgrade
sudo apt-get install python3-dev python3-pip libopenblas-base libopenmpi-dev libomp-dev libopenblas-dev libopenmpi-dev
sudo apt-get install libhdf5-serial-dev hdf5-tools libhdf5-dev zlib1g-dev zip libjpeg8-dev liblapack-dev libblas-dev gfortran

# Download a newer version of cmake for building PyTorch from source
wget https://github.com/Kitware/CMake/releases/download/v3.28.3/cmake-3.28.3-linux-aarch64.tar.gz
tar -xvzf cmake-3.28.3-linux-aarch64.tar.gz
export PATH=/path/to/cmake-3.28.3-linux-aarch64/bin:$PATH

# avalanche-lib requires this pytorch.distributed; so cannot use Jetson pre-built wheels
export USE_NCCL=0
export USE_DISTRIBUTED=1                    
export USE_QNNPACK=0
export USE_PYTORCH_QNNPACK=0
export TORCH_CUDA_ARCH_LIST="7.2;8.7"   # "7.2;8.7" for JetPack 5 wheels for Xavier/Orin
export PYTORCH_BUILD_VERSION=2.1.0  # without the leading 'v', e.g. 1.3.0 for PyTorch v1.3.0
export PYTORCH_BUILD_NUMBER=1

# Build GPU-enabled PyTorch from source for v2.1.0
git clone --recursive --branch v2.1.0 http://github.com/pytorch/pytorch
cd pytorch
pip install -r requirements.txt
pip install scikit-build
pip install ninja
python setup.py bdist_wheel
python setup.py install

TensorFlow refer to en or zh.