Zexin Li

Please keep honest, open, patient, happy and visionary.

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:
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    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)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
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.