在Jetson设备上安装PyTorch和Torchvision的完整指南

在Jetson设备上安装PyTorch和Torchvision的完整指南

前言

NVIDIA Jetson系列开发板（如Jetson Nano、Jetson TX2、Jetson Xavier NX等）是强大的边缘计算设备，广泛应用于计算机视觉、深度学习和AI项目中。PyTorch作为当前最流行的深度学习框架之一，在Jetson设备上的安装过程与常规PC有所不同。本文将详细介绍在Jetson设备上安装PyTorch及其配套库torchvision的完整步骤。

一、准备工作

1.1 确认Jetson设备信息

在开始安装前，首先需要确认你的Jetson设备型号和系统信息：

# 查看JetPack版本
cat /etc/nv_tegra_release

# 查看系统架构
uname -m

# 查看CUDA版本
nvcc --version

1.2 更新系统包

确保系统包是最新的：

sudo apt update
sudo apt upgrade -y
sudo apt autoremove -y

1.3 安装必要依赖

安装编译PyTorch所需的依赖项：

sudo apt install -y python3-pip libopenblas-base libopenmpi-dev libjpeg-dev zlib1g-dev

二、安装PyTorch

2.1 选择合适的PyTorch版本

由于Jetson设备基于ARM架构，不能直接使用PyTorch官方提供的pip安装包。NVIDIA为Jetson提供了预编译的PyTorch wheel文件。

根据你的JetPack版本选择合适的PyTorch版本：

可以在 https://developer.nvidia.com/embedded/downloads 下载到对应版本的torch。

2.2 下载并安装PyTorch

点开对应版本，下载whl后pip install即可。

2.3 验证PyTorch安装

import torch
print(torch.__version__)
print(torch.cuda.is_available())  # 应该返回True
print(torch.backends.cudnn.version())  # 查看cuDNN版本

如果此时报错 ImportError: libcusparseLt.so.0: cannot open shared object file: No such file or directory，说明缺少cusparselt库，去 https://developer.nvidia.com/cusparselt-downloads 安装即可，安装指令如下：

local:

wget https://developer.download.nvidia.com/compute/cusparselt/0.7.1/local_installers/cusparselt-local-tegra-repo-ubuntu2204-0.7.1_1.0-1_arm64.deb
sudo dpkg -i cusparselt-local-tegra-repo-ubuntu2204-0.7.1_1.0-1_arm64.deb
sudo cp /var/cusparselt-local-tegra-repo-ubuntu2204-0.7.1/cusparselt-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get -y install libcusparselt0 libcusparselt-dev

network:

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/arm64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt-get update
sudo apt-get -y install libcusparselt0 libcusparselt-dev

三、安装Torchvision

3.1 选择兼容的Torchvision版本

Torchvision版本必须与PyTorch版本兼容：

3.2 从源码编译安装Torchvision

由于没有预编译的torchvision wheel文件，我们需要从源码编译：

# 安装编译依赖
sudo apt install -y libjpeg-dev zlib1g-dev libpython3-dev libavcodec-dev libavformat-dev libswscale-dev

# 克隆torchvision源码（注意选择正确的版本分支）
git clone --branch v0.20.0 https://github.com/pytorch/vision torchvision

# 进入目录并安装
cd torchvision
export BUILD_VERSION=0.10.0
python3 setup.py install --user

# 返回上级目录并清理
cd ..

3.3 验证Torchvision安装

import torchvision
print(torchvision.__version__)

四、常见问题解决方案

4.1 安装过程中内存不足

Jetson设备内存有限，编译时可能出现内存不足的问题：

# 创建交换空间
sudo fallocate -l 4G /swapfile
sudo chmod 600 /swapfile
sudo mkswap /swapfile
sudo swapon /swapfile

# 永久添加交换空间（可选）
echo '/swapfile none swap sw 0 0' | sudo tee -a /etc/fstab

4.2 导入torch时出现Segmentation fault

这通常是由于版本不匹配造成的，请确保：

1. PyTorch、Torchvision和JetPack版本兼容
2. 使用正确的Python版本（通常为Python 3.6）
3. 彻底卸载后重新安装

4.3 CUDA不可用

如果torch.cuda.is_available()返回False：

# 检查CUDA是否安装正确
which nvcc

# 检查环境变量
echo $LD_LIBRARY_PATH

五、优化建议

5.1 使用虚拟环境

建议使用virtualenv或conda创建独立环境：

sudo apt install -y python3-venv
python3 -m venv pytorch_env
source pytorch_env/bin/activate

5.2 安装其他有用的库

pip3 install matplotlib pillow scipy pandas

5.3 性能优化

# 在代码中添加以下内容以获得最佳性能
torch.backends.cudnn.benchmark = True

结语

在Jetson设备上安装PyTorch和Torchvision虽然比在普通PC上复杂一些，但按照上述步骤操作应该能够顺利完成。安装完成后，你就可以在Jetson设备上运行各种深度学习模型了。如果在安装过程中遇到任何问题，可以参考NVIDIA官方论坛或PyTorch社区寻求帮助。