NVIDIA 优化框架：Jetson 平台 PyTorch 安装指南

NVIDIA 优化框架：Jetson 平台 PyTorch 安装指南

摘要

本文档提供了在 Jetson 平台安装 PyTorch 的操作说明。

1. 概述

1.1 Jetson 平台上的 PyTorch

适用于 JetPack 的 PyTorch 是一款经过优化的张量库，可用于深度学习，支持 GPU 和 CPU 运算。其通过基于tape-based的系统，在函数层和神经网络层实现自动微分。该功能为深度学习框架带来了高度的灵活性与速度，同时提供了类似 NumPy 但经过加速的功能。

NVIDIA 提供的这些可再分发软件是 PyTorch 的 Python pip wheel 安装程序，具备 GPU 加速能力且支持 cuDNN。这些包需按照文档说明，安装在指定版本的 JetPack 之上。

1.2 支持的 Jetson 设备

• Jetson AGX Xavier：全球首款用于自主机器的 AI 计算机，在功耗低于 30W 的嵌入式模块中，可提供 GPU 工作站级别的性能。
• Jetson AGX Orin：包含高性能、高功耗效率的 Jetson AGX Orin 模块，还可模拟其他 Jetson 模块。在与 Jetson AGX Xavier 相同的紧凑尺寸下，性能提升 8 倍，算力最高可达 275 TOPS，适用于开发高级机器人及其他自主机器产品。
• Jetson Xavier NX：以小型系统级模块（SoM）形态，将超级计算机性能带到边缘端。最高 21 TOPS 的加速计算能力，可并行运行现代神经网络，并处理来自多个高分辨率传感器的数据，满足完整 AI 系统的需求。

1.3 Jetson 平台 PyTorch 的优势

在 Jetson 平台安装 PyTorch，可让用户在轻量级移动平台上使用该框架的最新版本。

2. 前提条件与安装步骤

在 Jetson 设备上安装 PyTorch 前，请确保完成以下操作：

1. 在 Jetson 设备上安装 JetPack。

2. 安装 PyTorch 所需的系统包，执行以下命令：

   sudo apt-get -y update; 
   sudo apt-get install -y python3-pip libopenblas-dev;
   

3. 若安装的是 24.06 及后续版本的 PyTorch，需先安装 cusparselt，步骤如下：

   wget raw.githubusercontent.com/pytorch/pytorch/5c6af2b583709f6176898c/common/install_cusparselt.sh
   export CUDA_VERSION=12.1  # 此处以 12.1 为例
   bash ./install_cusparselt.sh
   

完成上述前提操作后，按以下步骤安装 PyTorch：

1. 定义环境变量（需根据兼容性矩阵替换对应版本号）：

   • JP_VERSION：JetPack 版本号

   • PYT_VERSION：PyTorch wheel 包的已发布版本号

     • Jetson Download Center | NVIDIA Developer
     • 

   • 设置 PyTorch 安装路径，示例如下：（可以不管）

export TORCH_INSTALL=https://developer.download.nvidia.cn/compute/
# 或本地路径（若已下载安装包）
export TORCH_INSTALL=path/to/torch-2.2.0a0+81ea7a4+nv23.12-cp38-cp38
```

2. 执行安装命令（通过官方链接安装示例，较慢建议先下载再安装）：

   pip3 install --no-cache https://developer.download.nvidia.com/compute/redist/jp/v$JP_VERSION/py$PYT_VERSION
   

2.1 安装多个 PyTorch 版本

若需同时保留多个 PyTorch 版本，可通过虚拟环境实现，步骤如下：

1. 创建虚拟环境：

   先安装 virtualenv 包，再创建 Python 3 虚拟环境：

   # 创建虚拟环境
   python -m venv myenv
   

2. 激活虚拟环境

      # Linux/Mac:
   source myenv/bin/activate# Windows:
   myenv\Scripts\activate
   

3. 安装指定版本 PyTorch（建议先把torch包下载到本地，再安装）

   # 在虚拟环境中安装，无需管理员权限
   pip3 install torch-2.0.0+nv23.05-cp38-cp38-linux_aarch64.whl
   

   

4. 退出虚拟环境

   deactivate
   

5. 运行指定版本 PyTorch：

   虚拟环境创建完成后，只需激活即可使用对应的 PyTorch 版本，使用完毕后退出环境：

   source myenv/bin/activate
   <执行所需的 PyTorch 脚本>
   deactivate
   

2.2 升级 PyTorch

若需升级到更新版本的 PyTorch（需确认新版本已发布），需先卸载当前版本，再参考 “前提条件与安装步骤” 安装目标版本，卸载命令如下：

sudo pip3 uninstall -y torch

3. 验证安装结果

操作目的

确认 PyTorch 已在 Jetson 平台成功安装，需启动 Python 终端并导入 PyTorch 库。

操作步骤

1. 在终端执行以下命令，配置环境变量并启动 Python 3：

   source myenv/bin/activate
   python3
   

2. 在 Python 交互环境中导入 PyTorch：

   >>> import torch
   

若上述命令执行无报错，则说明 PyTorch 安装成功。

4. 卸载 PyTorch

通过以下命令可轻松卸载 PyTorch：

sudo pip3 uninstall -y torch

5. 故障排除

若遇到 Jetson 平台相关问题，可加入 “NVIDIA Jetson 与嵌入式系统社区”（NVIDIA Jetson and Embedded Systems community）进行讨论。