【踩坑记录】NVIDIA Jetson (JetPack 6) 极简配置 PyTorch GPU 环境指南

梦醒过后说珍重
2 阅读4分钟

为什么学这个？

起初，我以为在 Jetson 上配环境和在普通的 x86 主机上一样，直接 pip install torch 加上 cu121 后缀就能搞定。结果现实狠狠给我上了一课：Jetson 是 ARM64 架构（aarch64），且底层调用 GPU 依赖于专属的 Tegra 芯片驱动方案。 普通的 PC 端安装包在这里完全行不通，默认安装只会得到一个纯 CPU 版本的 PyTorch。而且由于网络环境限制，直接使用在线源经常遭遇 SSL 中断。为了彻底打通这套边缘侧的加速计算环境，我经过多次试错，跑通了一套 100% 成功的离线配置+源码编译流程。在此记录，防患未然。

核心内容 / 步骤

在 Jetson 上配置环境，核心原则只有一个：严格匹配 JetPack 版本。

1. 摸清底细：确认 JetPack 与 L4T 版本

在终端输入以下命令查看底层系统版本：

Bash

cat /etc/nv_tegra_release

如果输出类似 R36 (release)...，说明使用的是最新的 JetPack 6.x 系列（Ubuntu 22.04 + Python 3.10，默认支持 CUDA 12.2）。

2. 环境隔离：创建虚拟环境

不要弄脏系统环境，使用 venv 隔离依赖：

Bash

python3 -m venv .venv
source .venv/bin/activate

3. 核心安装：使用 wget 离线下载并安装 PyTorch

我们可以通过增加 pip 的超时时间限制，并信任该域名来绕过一些严格的 SSL 拦截。

pip install torch torchvision torchaudio --index-url https://pypi.jetson-ai-lab.dev/jp6/cu122 --default-timeout=2000 --trusted-host pypi.jetson-ai-lab.dev

由于在线 pip 极易因为网络波动导致下载失败，若上述方法失败后，我们直接用 wget -c（支持断点续传）把 NVIDIA 官方编译好的大文件拉到本地。以 PyTorch 2.5.0 为例：

Bash

# 下载专属 wheel 包
wget -c https://developer.download.nvidia.com/compute/redist/jp/v61/pytorch/torch-2.5.0a0+872d972e41.nv24.08.17622132-cp310-cp310-linux_aarch64.whl

# 本地离线安装
pip install ./torch-2.5.0a0+872d972e41.nv24.08.17622132-cp310-cp310-linux_aarch64.whl

4. 特殊指令：源码编译 Torchvision

这是最容易卡壳的一步。在 Jetson 的离线环境下，torchvision 绝不能随便找个包安装，必须从源码编译，以确保其 C++ 扩展能和刚刚装好的专属版 PyTorch 完美链接。PyTorch 2.5.0 对应的 Torchvision 版本为 v0.20.0。

Bash

# 安装基础 C++ 图像处理库
sudo apt-get update
sudo apt-get install -y libjpeg-dev zlib1g-dev libpython3-dev libopenblas-dev libavcodec-dev libavformat-dev libswscale-dev

# 克隆对应版本的源码
git clone --branch v0.20.0 https://github.com/pytorch/vision torchvision
cd torchvision

# 声明版本环境变量，并启动编译安装（耗时约 5-15 分钟）
export BUILD_VERSION=0.20.0
python setup.py install

5. 验证安装

进入 Python 交互环境：

Python

import torch
import torchvision

print("PyTorch 版本:", torch.__version__)
print("Torchvision 版本:", torchvision.__version__)
print("GPU 是否可用:", torch.cuda.is_available())

if torch.cuda.is_available():
    print("显卡名称:", torch.cuda.get_device_name(0))

看到 True 和具体的显卡型号，配置大功告成！

遇到的问题与解决方法

  • 坑一：装完发现是 CPU 版本 (is_available() 为 False)

    • 原因： 习惯性地直接输入了 pip install torch。默认 PyPI 仓库检测到 ARM64 架构，直接分配了纯 CPU 包。
    • 解决： 果断 pip uninstall 卸载干净，必须去 NVIDIA 开发者网站获取针对 L4T 编译的专属 .whl 包。

  • 坑二：下载频繁报错 [SSL: UNEXPECTED_EOF_WHILE_READING]

    • 原因： 安装包太大（超 1GB），网络稍有波动 pip 就会强制中断。
    • 解决： 放弃 pip 直连，改为 wget -c 暴力断点续传下载到本地再安装。

  • 坑三：Torchvision 找不到匹配的 GPU 版本

    • 原因： 边缘设备的预编译生态不如 PC 完善，版本绑定极其严格。
    • 解决： 查阅版本对应表，下载确切版本的源码，务必通过 export BUILD_VERSION=x.xx.x 注入版本号后进行本地编译。

收获与总结

  1. 硬件架构决定软件生态： 边缘计算设备（ARM）和常规 PC（x86）完全是两套逻辑，切忌生搬硬套 PC 端的经验。
  2. 底层依赖的敬畏之心： 涉及到计算机视觉的任务，torchvision 的底层强依赖系统的 C++ 图像处理库（如 libjpeg-dev），编译前一定要检查系统依赖是否齐全。
  3. 稳字当头： 面对动辄几个 G 的环境包，wget -c 离线下载 + 本地源码编译，虽然看似繁琐，但往往是玄学网络环境下最节省时间的终极解法。
avatar
文章
阅读
粉丝