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

安裝方法

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

• 一種是直接安裝他人已經(jīng)編譯好的 PyTorch 輪子；
• 一種是自己從頭開始開始構(gòu)建 PyTorch 輪子并且安裝。

使用輪子安裝

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

安裝前先確保 python 版本為 PyTorch 輪子對應(yīng)的 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 倉庫 尋找。

編譯源碼安裝

接下來說一下重點(diǎn)，自己從源碼構(gòu)建輪子并安裝。自行構(gòu)建需要花費(fèi)數(shù)個小時的時間。

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后的版本參數(shù)。

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

擴(kuò)大交換內(nèi)存

構(gòu)建 PyTorch 需要大于 4GB 的 RAM 和 2GB 的交換內(nèi)存。

先使用free -m查看當(dāng)前的交換內(nèi)存。

安裝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

保存修改后重啟系統(tǒng)。

sudo reboot

重啟完畢后可以使用 free -m 查看交換內(nèi)存是否擴(kuò)大成功。

使用 GCC

請使用 GCC 和 G++ 編譯器，如果使用 GNU 編譯器會導(dǎo)致浮點(diǎn)數(shù)出現(xiàn)錯誤的問題，使用 Clang 編譯器在構(gòu)建 CUDA 11.4 的 PyTorch 時候會報(bào)錯。

設(shè)置 CUDA 版本

修改環(huán)境變量

修改系統(tǒng)環(huán)境變量，如果你使用的是 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 的環(huán)境變量，則直接在對應(yīng)行修改。

保存并退出，應(yīng)用修改：

source ~/.zshrc #或 ~/.bashrc

切換 CUDA 版本

先切換系統(tǒng) 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)

保存并退出。

構(gòu)建 PyTorch

設(shè)置 Ninja 構(gòu)建參數(shù)

需要添加一些環(huán)境變量。

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/cuda
sudo ln -s /usr/lib/aarch64-linux-gnu/libcublas.so /usr/local/cuda/lib64/libcublas.so

開始構(gòu)建 PyTorch

python3 setup.py bdist_wheel

如果構(gòu)建過程中報(bào)錯，需要重新構(gòu)建，請先用python3 setup.py clean重置構(gòu)建。

安裝 PyTorch

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

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

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

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

驗(yàn)證安裝

導(dǎo)入驗(yàn)證 torch 不能在 PyTorch 源代碼目錄，否則會報(bào)錯，先用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 辛辛苦苦構(gòu)建 PyTorch 輪子備份保存一下。

mv pytorch/dist/*.whl ~

刪除源代碼：

rm -r pytorch

總結(jié)

以上為個人經(jīng)驗(yàn)，希望能給大家一個參考，也希望大家多多支持腳本之家。

最新評論