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())