Jetson Xavier NX安装CUDA支持的PyTorch教程

 更新时间:2026年04月09日 08:52:58   作者:jax在coding  
本文介绍了在Jetson Xavier NX上安装CUDA支持的PyTorch的方法,提供了两种安装方式,一种是直接安装编译好的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

总结

以上为个人经验,希望能给大家一个参考,也希望大家多多支持脚本之家。

相关文章

  • Python之多线程退出与停止的一种实现思路

    Python之多线程退出与停止的一种实现思路

    这篇文章主要介绍了Python之多线程退出与停止的一种实现思路,具有很好的参考价值,希望对大家有所帮助,如有错误或未考虑完全的地方,望不吝赐教
    2024-02-02
  • 完美解决Python 2.7不能正常使用pip install的问题

    完美解决Python 2.7不能正常使用pip install的问题

    今天小编就为大家分享一篇完美解决Python 2.7不能正常使用pip install的问题,具有很好的参考价值,希望对大家有所帮助。一起跟随小编过来看看吧
    2018-06-06
  • pyinstaller打包遇到dll问题及解决方案

    pyinstaller打包遇到dll问题及解决方案

    文章讲述了解决xgboost.dll问题的两种方法:步骤一是在虚拟环境中将xgboost包的文件夹复制到项目路径,并在spec配置文件中声明打包进exe文件;步骤二是在Anaconda的base环境中安装xgboost,以备步骤一失败时使用
    2026-01-01
  • Python库安装加速之使用清华大学镜像源的详细指南

    Python库安装加速之使用清华大学镜像源的详细指南

    这篇文章主要介绍了如何通过清华大学镜像源加速Python库安装,提供临时指定和永久配置方法,并列举阿里云、中科大、豆瓣等其他镜像源,文中通过代码介绍的非常详细,需要的朋友可以参考下
    2025-06-06
  • windowns使用PySpark环境配置和基本操作

    windowns使用PySpark环境配置和基本操作

    pyspark是Spark对Python的api接口,可以在Python环境中通过调用pyspark模块来操作spark,这篇文章主要介绍了windowns使用PySpark环境配置和基本操作,感兴趣的可以了解一下
    2021-05-05
  • python中使用pyhook实现键盘监控的例子

    python中使用pyhook实现键盘监控的例子

    这篇文章主要介绍了python中使用pyhook实现键盘监控的例子,包含pyhook的下载地址和手册地址及一个Windows下的监控实例,需要的朋友可以参考下
    2014-07-07
  • 如何实现删除numpy.array中的行或列

    如何实现删除numpy.array中的行或列

    如何实现删除numpy.array中的行或列?今天小编就为大家分享一篇对删除numpy.array中行或列的实例讲解,具有很好的参考价值,希望对大家有所帮助。一起跟随小编过来看看吧
    2018-05-05
  • Python去掉字符串中空格的方法

    Python去掉字符串中空格的方法

    这篇文章主要介绍了Python中去掉字符串中空格的方法,使用了strip()、lstrip()、rstrip()函数,需要的朋友可以参考下
    2014-03-03
  • 在树莓派2或树莓派B+上安装Python和OpenCV的教程

    在树莓派2或树莓派B+上安装Python和OpenCV的教程

    这篇文章主要介绍了在树莓派2或树莓派B+上安装Python和OpenCV的教程,主要基于GTK库,并以Python2.7和OpenCV 2.4.X版本的安装作为示例,需要的朋友可以参考下
    2015-03-03
  • 详解Python中的上下文管理器原理

    详解Python中的上下文管理器原理

    这篇文章主要为大家详细介绍了Python中的上下文管理器的原理与使用,文中的示例代码讲解详细,感兴趣的小伙伴可以跟随小编一起学习一下
    2023-03-03

最新评论