PyTorch for Jetson

Autonomous Machines Jetson & Embedded Systems Announcements
, ,
1

Below are pre-built PyTorch pip wheel installers for Jetson Nano, TX1/TX2, Xavier, and Orin with JetPack 4.2 and newer.

Download one of the PyTorch binaries from below for your version of JetPack, and see the installation instructions to run on your Jetson. These pip wheels are built for ARM aarch64 architecture, so run these commands on your Jetson (not on a host PC). You can also use the containers from jetson-containers.

PyTorch pip wheels

JetPack 6

PyTorch v2.3.0
PyTorch v2.2.0
PyTorch v2.1.0

JetPack 5

PyTorch v2.1.0
PyTorch v2.0.0
PyTorch v1.14.0
PyTorch v1.13.0
PyTorch v1.12.0
PyTorch v1.11.0

JetPack 4

PyTorch v1.10.0
  • JetPack 4.4 (L4T R32.4.3) / JetPack 4.4.1 (L4T R32.4.4) / JetPack 4.5 (L4T R32.5.0) / JetPack 4.5.1 (L4T R32.5.1) / JetPack 4.6 (L4T R32.6.1)
PyTorch v1.9.0
PyTorch v1.8.0
PyTorch v1.7.0
PyTorch v1.6.0
  • JetPack 4.4 (L4T R32.4.3) / JetPack 4.4.1 (L4T R32.4.4) / JetPack 4.5 (L4T R32.5.0) / JetPack 4.5.1 (L4T R32.5.1) / JetPack 4.6 (L4T R32.6.1)
    • Python 3.6 - torch-1.6.0-cp36-cp36m-linux_aarch64.whl
    • The JetPack 4.4 production release (L4T R32.4.3) only supports PyTorch 1.6.0 or newer, due to updates in cuDNN.
    • This wheel of the PyTorch 1.6.0 final release replaces the previous wheel of PyTorch 1.6.0-rc2.
PyTorch v1.5.0
PyTorch v1.4.0
PyTorch v1.3.0
PyTorch v1.2.0
PyTorch v1.1.0
PyTorch v1.0.0

Instructions

Installation

Below are example commands for installing these PyTorch wheels on Jetson. Substitute the URL and filenames from the desired PyTorch download from above.

> Python 3

# substitute the link URL and wheel filename from the desired torch version above
wget https://nvidia.box.com/shared/static/p57jwntv436lfrd78inwl7iml6p13fzh.whl -O torch-1.8.0-cp36-cp36m-linux_aarch64.whl
sudo apt-get install python3-pip libopenblas-base libopenmpi-dev libomp-dev
pip3 install 'Cython<3'
pip3 install numpy torch-1.8.0-cp36-cp36m-linux_aarch64.whl

> Python 2.7

# substitute the link URL and wheel filename from the desired torch version above
wget https://nvidia.box.com/shared/static/1v2cc4ro6zvsbu0p8h6qcuaqco1qcsif.whl -O torch-1.4.0-cp27-cp27mu-linux_aarch64.whl
sudo apt-get install libopenblas-base libopenmpi-dev   # skip libopenmpi-dev for PyTorch >= 1.12
pip install future torch-1.4.0-cp27-cp27mu-linux_aarch64.whl

(PyTorch v1.4.0 for L4T R32.4.2 is the last version to support Python 2.7)

> torchvision

$ sudo apt-get install libjpeg-dev zlib1g-dev libpython3-dev libopenblas-dev libavcodec-dev libavformat-dev libswscale-dev
$ git clone --branch <version> https://github.com/pytorch/vision torchvision   # see below for version of torchvision to download
$ cd torchvision
$ export BUILD_VERSION=0.x.0  # where 0.x.0 is the torchvision version  
$ python3 setup.py install --user
$ cd ../  # attempting to load torchvision from build dir will result in import error
$ pip install 'pillow<7' # always needed for Python 2.7, not needed torchvision v0.5.0+ with Python 3.6

Select the version of torchvision to download depending on the version of PyTorch that you have installed:

  • PyTorch v1.0 - torchvision v0.2.2
  • PyTorch v1.1 - torchvision v0.3.0
  • PyTorch v1.2 - torchvision v0.4.0
  • PyTorch v1.3 - torchvision v0.4.2
  • PyTorch v1.4 - torchvision v0.5.0
  • PyTorch v1.5 - torchvision v0.6.0
  • PyTorch v1.6 - torchvision v0.7.0
  • PyTorch v1.7 - torchvision v0.8.1
  • PyTorch v1.8 - torchvision v0.9.0
  • PyTorch v1.9 - torchvision v0.10.0
  • PyTorch v1.10 - torchvision v0.11.1
  • PyTorch v1.11 - torchvision v0.12.0
  • PyTorch v1.12 - torchvision v0.13.0
  • PyTorch v1.13 - torchvision v0.13.0
  • PyTorch v1.14 - torchvision v0.14.1
  • PyTorch v2.0 - torchvision v0.15.1
  • PyTorch v2.1 - torchvision v0.16.1
  • PyTorch v2.2 - torchvision v0.17.1
  • PyTorch v2.3 - torchvision v0.18.0

Verification

To verify that PyTorch has been installed correctly on your system, launch an interactive Python interpreter from terminal (python command for Python 2.7 or python3 for Python 3.6) and run the following commands:

>>> import torch
>>> print(torch.__version__)
>>> print('CUDA available: ' + str(torch.cuda.is_available()))
>>> print('cuDNN version: ' + str(torch.backends.cudnn.version()))
>>> a = torch.cuda.FloatTensor(2).zero_()
>>> print('Tensor a = ' + str(a))
>>> b = torch.randn(2).cuda()
>>> print('Tensor b = ' + str(b))
>>> c = a + b
>>> print('Tensor c = ' + str(c))

>>> import torchvision
>>> print(torchvision.__version__)
Build from Source

Below are the steps used to build the PyTorch wheels. These were compiled in a couple of hours on a Xavier for Nano, TX2, and Xavier.

Note that if you are trying to build on Nano, you will need to mount a swap file.

Max Performance

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

Download PyTorch sources

$ git clone --recursive --branch <version> http://github.com/pytorch/pytorch
$ cd pytorch

Apply Patch
Select the patch to apply from below based on the version of JetPack you’re building on. The patches avoid the “too many CUDA resources requested for launch” error (PyTorch issue #8103, in addition to some version-specific bug fixes.

If you are applying one of the above patches to a different version of PyTorch, the file line locations may have changed, so it is recommended to apply these changes by hand.

Set Build Options

$ 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="5.3;6.2;7.2"   # or "7.2;8.7" for JetPack 5 wheels for Xavier/Orin

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

(remember to re-export these environment variables if you change terminal)

Build wheel for Python 2.7 (to pytorch/dist)

$ sudo apt-get install python-pip cmake libopenblas-dev libopenmpi-dev 
$ pip install -U pip

$ sudo pip install -U setuptools
$ sudo pip install -r requirements.txt

$ pip install scikit-build --user
$ pip install ninja --user

$ python setup.py bdist_wheel

Build wheel for Python 3.6 (to pytorch/dist)

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

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

$ python3 setup.py bdist_wheel
Note on Upgrading pip

If you get this error from pip/pip3 after upgrading pip with “pip install -U pip”:

pip
Traceback (most recent call last):
  File "/usr/bin/pip", line 9, in <module>
    from pip import main
ImportError
: cannot import name 'main'

You can either downgrade pip to it’s original version:

# Python 2.7
$ sudo python -m pip uninstall pip && sudo apt install python-pip --reinstall

# Python 3.6
$ sudo python3 -m pip uninstall pip && sudo apt install python3-pip --reinstall

-or- you can patch /usr/bin/pip (or /usr/bin/pip3)

diff --git a/pip b/pip
index 56bbb2b..62f26b9 100755
--- a/pip
+++ b/pip
@@ -6,6 +6,6 @@ import sys
 # Run the main entry point, similarly to how setuptools does it, but because
 # we didn't install the actual entry point from setup.py, don't use the
 # pkg_resources API.
-from pip import main
+from pip import __main__
 if __name__ == '__main__':
-    sys.exit(main())
+    sys.exit(__main__._main())
79 Likes
2

I met a trouble on installing Pytorch.
The Numpy module need python3-dev, but I can’t find ARM python3-dev for Python3.6. The source only includes the ARM python3-dev for Python3.5.1-3.
So, can the Python3.6 Pytorch work on Python3.5? Or where can I find ARM python3-dev for Python3.6 which is needed to install numpy?
Thanks a lot

5 Likes
3

Hmm that’s strange, on my system ‘sudo apt-get install python3-dev’ shows python3-dev version 3.6.7-1~18.04 is installed. Does it help if you run ‘sudo apt-get update’ before?

2 Likes
4

If I may ask, is there any way we could get the binaries for the Pytorch C++ frontend? I’m trying to build it from source and that would be really nice.

Thanks!

1 Like
5

I’m not sure that these are included in the distributable wheel since that’s intended for Python - so you may need to build following the instructions above, but with “python setup.py develop” or “python setup.py install” in the final step (see here).

6

@dusty_nv theres a small typo in the verification example, you’ll want to “import torch” not “pytorch”

Could be worth adding the “pip3 install numpy” into the steps, it worked for me first time, I didn’t hit the problem @buptwlr did with python3-dev being missing.

For python2 I had to “pip install future” before I could import torch (was complaining with “ImportError: No module named builtins”), apart from that it looks like its working as intended.

Thanks !!

1 Like
7

Woops, thanks for pointing that out Balnog, I have fixed that in the steps above.

Do you think that is only needed if you are building from source, or do you need to explicitly install numpy even if just using the wheel?

1 Like
8

I had to install numpy when using the python3 wheel. It wasn’t necessary for python2 - I’d not installed anything other than pip/pip3 on the system at that point (using the latest SD card image).

9

@dusty_nv , @Balnog
I have found the solution. I changed the apt source for speed reason.
And after putting the original sources back to the sources.list file, I successfully find the apt package.

10

OK thanks, I updated the pip3 install instructions to include numpy in case other users have this issue.

1 Like
11

from china
burn in jetson-nano-sd-r32.1-2019-03-18.img today.
download torch-1.1.0a0+b457266-cp36-cp36m-linux_aarch64.whl

-bash: pip3: command not found

do I need to install pip3?

sudo apt-get update failed.

python3-pip or python3-dev can’t be located.

edit /etc/apt/source.list to Chinese images failed again.

help!

12

Hi huhai, if apt-get update failed, that would prevent you from installing more packages from Ubuntu repo. Are you behind a firewall that is preventing you from connecting to the Ubuntu package repositories?

13

pip3 installed using:
https://bootstrap.pypa.io/get-pip.py
referencing :
https://packaging.python.org/tutorials/installing-packages/#ensure-you-can-run-pip-from-the-command-line

pip3 install numpy stlll failed.

restored /etc/apt/sources.list
apt-get work fine. Its the network , should be.
reinstalled pip3

numpy installed ok using:
pip3 install numpy --user

Playing ubuntu 16.04 and pytorch on this network for a while already, apt-get works well before.

New to ubuntu 18.04 and arm port, will keep working on apt-get .

turn out the wheel file can’t be download from china.
the file downloaded before have zero byte.
using an aliyun esc in usa finished the download job.

finally pytorch installed.

torch.cuda.is_available()
True

PS: compiling pytorch using jetson nano is a nightmare .

14

@dusty_nv ,
Hi, could you tell me how to install torchvision?
I cant install it by “pip3 install torchvision” cause it would collecting torch(from torchvision), and PyTorch does not currently provide packages for PyPI.

Please help me out. Thanks a lot

15

Hi buptwlr, run the commands below to install torchvision. It is installed from source:

$ git clone https://github.com/pytorch/vision
$ cd vision
$ sudo python setup.py install
2 Likes
16

When installing torchvision, I found I needed to install libjpeg-dev (using sudo apt-get install libjpeg-dev) becaue it’s required by Pillow which in turn is required by torchvision.

1 Like
17

Hi,
I’m getting a weird error while importing. I installed using the pre-built wheel specified in the top post.

>>> import torch
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/nvidia/python3/lib/python3.6/site-packages/torch/__init__.py", line 97, in <module>
    from torch._C import *
ImportError: libcusparse.so.10.0: cannot open shared object file: No such file or directory

I’m using a Xavier with the following CUDA version.

>>> cat /usr/local/cuda-10.0/version.txt 
CUDA Version 10.0.117
2 Likes
18

Hi haranbolt, have you re-flashed your Xavier with JetPack 4.2? Are you able to find cusparse library?

19

I had flashed it using JetPack 4.1.1 Developer preview. Is it necessary to reflash it using JetPack 4.2?

20

Yes, these PyTorch pip wheels were built against JetPack 4.2.