Jetson Xavier NX 安装 CUDA 支持的 PyTorch 指南

本指南将帮助开发者完成在 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/cuda

sudo 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