在 NVIDIA Orin (JetPack 6.0) 上安装 PyTorch 2.4 + Torchvision 0.19

在 NVIDIA Orin (JetPack 6.0) 上安装 PyTorch 2.4 + Torchvision 0.19 (含 CUDA 支持)

教程：在 NVIDIA Orin (JetPack 6.0) 上安装 PyTorch 2.4 + Torchvision 0.19 (含 CUDA 支持)

背景:

目标平台: NVIDIA Jetson Orin 系列设备
操作系统: Ubuntu 22.04 (aarch64)
JetPack 版本: JetPack 6.0 (L4T R36.x)
CUDA 版本: 12.2
目标: 在 Miniconda 环境 (Python 3.10) 中安装兼容的、支持 CUDA 的 PyTorch 和 Torchvision，用于运行多模态等深度学习模型。

遇到的主要困难 (踩坑记录):

在尝试为 JetPack 6.0 安装最新的 PyTorch (特别是 2.4.0a0 夜间构建版) 时，我们遇到了以下典型问题：

Torchvision 缺失: NVIDIA 官方为 JP6 提供的 PyTorch 索引 (https://developer.download.nvidia.com/compute/redist/jp/v60/pytorch/) 中，常常只包含 torch 的 wheel 文件，而缺少配套的 torchvision 和 torchaudio 文件。这导致无法通过简单的 pip install torchvision --extra-index-url … 来安装。
版本不兼容导致运行时错误: 即便尝试从源码编译 torchvision (例如 v0.19.0 版本)，如果 PyTorch 版本 (如 nv24.07 或 nv24.05 构建版) 与之不完全匹配，会在导入 torchvision 时出现运行时错误，例如我们遇到的 RuntimeError: operator torchvision::nms does not exist，或者论坛中其他人遇到的 AttributeError: … ‘register_fake’、‘set_python_module’ 等。这表明库内部的 C++/CUDA 接口不匹配。
CPU-Only 安装: 在尝试安装看似版本号匹配的组合 (如 PyTorch 2.3.1 + Torchvision 0.18.1) 时，如果 pip 没有正确地从 NVIDIA 源获取到 CUDA 版本的 wheel，可能会错误地安装成 CPU-only 的 PyTorch，导致 torch.cuda.is_available() 返回 False。

最终成功的解决方案 (基于 NVIDIA 论坛官方人员指导): 参考论坛链接： https://forums.developer.nvidia.com/t/compiling-torchvision-0-19-0-for-torch-2-4-0a0-07cecf4168-nv24-05-14710581/306300

经过多次尝试和参考 NVIDIA 开发者论坛 (参考帖子链接)，最终发现一个可行的组合是安装特定版本的 PyTorch 夜间构建版 (nv24.06)，并配合特定依赖 (cuSPARSELt)，然后从源码编译特定 tag 的 torchvision (v0.19.0)。

详细安装步骤:

第 1 步：环境准备 (Miniconda + Python 3.10)

下载 Miniconda (aarch64):
Bash

wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-aarch64.sh
安装 Miniconda:
Bash

bash Miniconda3-latest-Linux-aarch64.sh
同意许可，确认路径，并在询问是否 conda init 时输入 yes。
关闭并重新打开终端，或运行 source ~/.bashrc。
创建并激活 Conda 环境: (我们使用 lkalert 作为环境名，你可以自定义)
Bash

conda create -n lkalert python=3.10 -y
conda activate lkalert
更新 pip:
Bash

pip install --upgrade pip
第 2 步：清理旧环境 (如果之前有安装尝试)

Bash

pip uninstall torch torchvision torchaudio -y

删除可能残留的源码文件夹

cd ~
rm -rf ~/Downloads/vision ~/Downloads/torchvision
第 3 步：安装特定 PyTorch 版本 (nv24.06)

使用 NVIDIA 论坛帖子中 AastaLLL 提供的精确 URL:

Bash

pip install --no-cache https://developer.download.nvidia.com/compute/redist/jp/v60/pytorch/torch-2.4.0a0+f70bd71a48.nv24.06.15634931-cp310-cp310-linux_aarch64.whl
第 4 步：安装关键依赖 cuSPARSELt

这个库对于某些操作的兼容性似乎很重要。

Bash

下载 cuSPARSELt 仓库定义文件

wget https://developer.download.nvidia.com/compute/cusparselt/0.6.2/local_installers/cusparselt-local-tegra-repo-ubuntu2204-0.6.2_1.0-1_arm64.deb

安装仓库定义

sudo dpkg -i cusparselt-local-tegra-repo-ubuntu2204-0.6.2_1.0-1_arm64.deb

复制密钥环

sudo cp /var/cusparselt-local-tegra-repo-ubuntu2204-0.6.2/cusparselt-*-keyring.gpg /usr/share/keyrings/

更新 apt 列表

sudo apt-get update

安装 cuSPARSELt 库和开发文件

sudo apt-get install -y libcusparselt0 libcusparselt-dev

(可选) 清理下载的 .deb 文件

rm cusparselt-local-tegra-repo-ubuntu2204-0.6.2_1.0-1_arm64.deb
第 5 步：安装 Torchvision 构建依赖

Bash

sudo apt-get install -y cmake libjpeg-dev zlib1g-dev libpng-dev build-essential git
第 6 步：克隆、检出并构建 Torchvision v0.19.0

严格按照 AastaLLL 的步骤，特别是 git checkout 和 setup.py install 命令。

克隆仓库：
Bash

cd ~/Downloads # 或者你选择的其他路径
git clone https://github.com/pytorch/vision.git torchvision
cd torchvision/
检出 v0.19.0 标签：
Bash

git checkout v0.19.0
构建并安装 (使用论坛命令)：
Bash

python3 setup.py install --user
注意：–user 会将包安装到 ~/.local/lib/python3.10/site-packages/。这通常也能在 Conda 环境中被找到并使用。
第 7 步：验证安装

确保不在源码目录中：
Bash

cd ~
运行验证脚本：
Bash

python -c $‘import torch\nimport torchvision\n\nprint(f"— Verification —“)\nprint(f"PyTorch Version: {torch.version}”)\nprint(f"Torchvision Version: {torchvision.version}“)\n\ncuda_available = torch.cuda.is_available()\nprint(f"CUDA Available: {cuda_available}”)\n\nif cuda_available:\n print(f"CUDA Version (PyTorch): {torch.version.cuda}“)\n try:\n gpu_name = torch.cuda.get_device_name(0)\n print(f"GPU Name: {gpu_name}”)\n # Perform a tiny CUDA operation\n tensor_gpu = torch.randn(1, device=“cuda:0”)\n print(f"CUDA device functional: True (Tensor created on {tensor_gpu.device})“)\n except Exception as e:\n print(f"CUDA device functional: False (Error: {e})”)\nelse:\n print(“CUDA not available, skipping device tests.”)\nprint(f"--------------------")’
成功的验证输出应类似于：

— Verification —
PyTorch Version: 2.4.0a0+f70bd71a48.nv24.06…
Torchvision Version: 0.19.0a0+…
CUDA Available: True
CUDA Version (PyTorch): 12.2
GPU Name: Orin
CUDA device functional: True (Tensor created on cuda:0)

总结:

在 Jetson Orin (JetPack 6.0) 上安装最新的 PyTorch 夜间构建版和对应的 Torchvision 可能需要精确的版本匹配和特定的依赖。通过参考 NVIDIA 官方论坛并严格遵循 AastaLLL 提供的针对 PyTorch nv24.06 构建版和 Torchvision v0.19.0 的步骤（包括安装 cuSPARSELt），我们最终成功配置了支持 CUDA 的环境。