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
git clone http://github.com/pytorch/pytorch
cd pytorch
# later similar to TensorFlow build (The specific procedure is omitted here)

TensorFlow refer to en or zh.