本指南將幫助開發者完成在 Jetson Xavier NX 上安裝 CUDA 支持的 PyTorch。

安裝方法

在 Jetson 上安裝 Pytorch 只有兩種方法。

• 一種是直接安裝他人已經編譯好的 PyTorch 輪子；
• 一種是自己從頭開始開始構建 PyTorch 輪子并且安裝。

使用輪子安裝

可以從我的?GitHub 倉庫?直接下載我編譯好的 PyTorch 輪子（torch-1.13.0-cuda-11.4-python-3.8-aarch64）。

安裝前先確保 python 版本為 PyTorch 輪子對應的 Python 3.8。

下載完畢后使用pip install安裝：

sudo -H pip install torch-1.13.0a0+git7c98e70-cp38-cp38-linux_aarch64.whl

如果你以前安裝過 torch 的其他版本，則需要添加--force-reinstall覆蓋安裝。

其他低于 CUDA 11 的 PyTorch 輪子可以在?Qengineering 的 GitHub 倉庫?尋找。

編譯源碼安裝

接下來說一下重點，自己從源碼構建輪子并安裝。自行構建需要花費數個小時的時間。

PyTroch 1.11.0 以及以上版本只能在 Ubuntu 20.04 上安裝。

安裝依賴

sudo apt-get update
sudo apt-get upgrade

sudo apt-get install ninja-build git cmake 
sudo apt-get install libjpeg-dev libopenmpi-dev libomp-dev ccache libopenblas-dev libblas-dev libeigen3-dev

sudo pip3 install -U --user wheel mock pillow

sudo -H pip3 install testresources setuptools==58.3.0 scikit-build

下載源碼

請克隆自己需要的 PyTorch 版本，修改-b后的版本參數。

git clone -b v1.13.0 --depth=1 --recursive https://github.com/pytorch/pytorch.git

cd pytorch

擴大交換內存

構建 PyTorch 需要大于 4GB 的 RAM 和 2GB 的交換內存。

先使用free -m查看當前的交換內存。

安裝dphys-swapfiel：

sudo apt-get install dphys-swapfile

需要修改兩個文件/sbin/dphys-swapfile和/etc/dphys-swapfile。

先修改/sbin/dphys-swapfile的第30行，將CONF_MAXSWAP=2048修改成4096：

sudo vim /sbin/dphys-swapfile

再修改/etc/dphys-swapfile的第26行，將#去掉，后面加上4096：

sudo vim /etc/dphys-swapfile

保存修改后重啟系統。

sudo reboot

重啟完畢后可以使用?free -m?查看交換內存是否擴大成功。

使用 GCC

請使用 GCC 和 G++ 編譯器，如果使用 GNU 編譯器會導致浮點數出現錯誤的問題，使用 Clang 編譯器在構建 CUDA 11.4 的 PyTorch 時候會報錯。

設置 CUDA 版本

修改環境變量

修改系統環境變量，如果你使用的是 bash 則修改?~/.bashrc，如果使用 zsh 則修改?~/.zshrc。

vim ~/.zshrc #或 ~/.bashrc

在文件末尾添加三行：

export CUDA_HOME=/usr/local/cuda
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64
export PATH=$PATH:$CUDA_HOME/bin

如果你之前添加過 CUDA 的環境變量，則直接在對應行修改。

保存并退出，應用修改：

source ~/.zshrc #或 ~/.bashrc

切換 CUDA 版本

先切換系統 CUDA 版本到即將編譯的 CUDA 版本，這里以 CUDA 11.4 為例。

sudo update-alternatives --config cuda

選擇2，回車。

修改 CMakeLists.txt

需要在 CMakeLists.txt 文件里指定 CUDA 編譯器的路徑。

cd pytorch
vim CMakeLists.txt

在?project(Torch CXX C)上方添加一行?set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc)

保存并退出。

構建 PyTorch

設置 Ninja 構建參數

需要添加一些環境變量。

export BUILD_CAFFE2_OPS=OFF
export USE_FBGEMM=OFF
export USE_FAKELOWP=OFF
export BUILD_TEST=OFF
export USE_MKLDNN=OFF
export USE_NNPACK=OFF
export USE_XNNPACK=OFF
export USE_QNNPACK=OFF
export USE_PYTORCH_QNNPACK=OFF
export USE_CUDA=ON
export USE_CUDNN=ON
export TORCH_CUDA_ARCH_LIST="5.3;6.2;7.2"
export USE_NCCL=OFF
export USE_SYSTEM_NCCL=OFF
export USE_OPENCV=OFF
export MAX_JOBS=4
export PATH=/usr/lib/ccache:$PATH
export CC=gcc
export CXX=g++
export CUDACXX=/usr/local/cuda/bin/nvcc
export CUDA_TOOLKIT_ROOT_DIR=/usr/local/cudasudo ln -s /usr/lib/aarch64-linux-gnu/libcublas.so /usr/local/cuda/lib64/libcublas.so

開始構建 PyTorch

python3 setup.py bdist_wheel

如果構建過程中報錯，需要重新構建，請先用python3 setup.py clean重置構建。

安裝 PyTorch

經過漫長的編譯，編譯好的.whl輪子文件會放在dist文件夾里。

cd dist
sudo -H pip3 install <torch文件>.whl

如果你以前安裝過 torch 的其他版本，則需要添加--force-reinstall覆蓋安裝：

sudo -H pip3 install <torch文件>.whl --force-reinstall

驗證安裝

導入驗證 torch 不能在 PyTorch 源代碼目錄，否則會報錯，先用cd ../..退出源代碼文件夾。

python3 -c "import torch as t; print(t.__version__); print(t.version.cuda); print(t.cuda.is_available());"

大功告成！

清理空間

先停 dphys-swapfile：

sudo /etc/init.d/dphys-swapfile stop

如果磁盤空間緊缺，可以將 PyTorch 源代碼刪除，不過刪除前建議先將 Jetson 辛辛苦苦構建 PyTorch 輪子備份保存一下。

mv pytorch/dist/*.whl ~

刪除源代碼：

rm -r pytorch