jetpack5.0配置对应版本的torch和 torchvision

打开官网：https://forums.developer.nvidia.com/t/pytorch-for-jetson/72048

找到对应的jetpack5 ，我选择了 PyTorch v1.12.0

找到对应的torchvision

我需要对应的torchvision v0.13.0（对应PyTorch 1.12.0）

清楚对应版本后，就是具体配置了

因为torch2.3.0才开始提供编译好的torchvision

所以我们需要自己去编译

Jetson Orin NX PyTorch 环境配置指南

环境信息

设备: Jetson Orin NX
系统: JetPack 5.0 (L4T R34.1)
Python: 3.8
CUDA: 11.4
虚拟环境: rtgs_venv

完整操作步骤

步骤1：准备工作

复制代码

# 1.1 创建并激活虚拟环境
python3 -m venv rtgs_venv
source rtgs_venv/bin/activate

# 1.2 安装基础依赖
sudo apt-get update
sudo apt-get install python3-pip libopenblas-base libopenmpi-dev libomp-dev
pip3 install 'Cython<3'
pip3 install numpy

步骤2：安装PyTorch

复制代码

# 2.1 下载PyTorch wheel文件（针对JetPack 5.0）
wget https://developer.download.nvidia.com/compute/redist/jp/v50/pytorch/torch-1.12.0a0+2c916ef.nv22.3-cp38-cp38-linux_aarch64.whl

# 2.2 安装PyTorch
pip3 install torch-1.12.0a0+2c916ef.nv22.3-cp38-cp38-linux_aarch64.whl

步骤3：安装torchvision

复制代码

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

# 3.2 下载torchvision源码（对应PyTorch 1.12.0）
# 方案A: 使用git克隆（需要稳定网络）
git clone --branch v0.13.0 https://github.com/pytorch/vision torchvision
cd torchvision

# 方案B: 如果git clone中断，可从GitHub下载zip压缩包
# 访问 https://github.com/pytorch/vision/releases/tag/v0.13.0
# 下载源码zip包，解压后重命名为torchvision

# 3.3 修正环境变量（关键步骤）
# 检查并修正CUDA_HOME设置
echo $CUDA_HOME  # 如果显示 :/usr/local/cuda（开头有冒号）
unset CUDA_HOME
export CUDA_HOME=/usr/local/cuda

# 3.4 编译安装（预计10-30分钟）
python3 setup.py install --user

编译结果：

步骤4：验证安装

复制代码

# 4.1 验证PyTorch
python -c "import torch; print('PyTorch版本:', torch.__version__); print('CUDA版本:', torch.version.cuda)"

# 4.2 验证torchvision
python -c "import torchvision; print('TorchVision版本:', torchvision.__version__)"

# 4.3 验证CUDA支持
python -c "import torch; print('CUDA是否可用:', torch.cuda.is_available())"

关键问题与解决方案

问题1：torchvision编译错误

错误信息:

复制代码

error: [Errno 2] No such file or directory: ':/usr/local/cuda/bin/nvcc'

原因 : CUDA_HOME环境变量被错误设置为: /usr/local/cuda（开头有冒号）

解决方案:

复制代码

# 修正CUDA_HOME
unset CUDA_HOME
export CUDA_HOME=/usr/local/cuda

# 验证修正
echo $CUDA_HOME  # 应该输出：/usr/local/cuda（无冒号开头）

问题2：git clone中断

解决方案 : 前往 https://github.com/pytorch/vision/releases/tag/v0.13.0 下载zip压缩包，然后拷贝到Jetson设备上

最终环境配置结果

✅ PyTorch: 1.12.0a0+2c916ef.nv22.3 (NVIDIA优化版)
✅ torchvision: 0.13.0 (源码编译)
✅ CUDA: 11.4 (支持GPU加速)
✅ Python: 3.8
✅ 虚拟环境: rtgs_venv (已激活)

使用建议

每次使用前激活虚拟环境

复制代码

source /home/jetson/aaa-rtgs/rtgs_venv/bin/activate

快速验证环境

复制代码

python -c "import torch; import torchvision; print(f'PyTorch: {torch.__version__}, TorchVision: {torchvision.__version__}, CUDA可用: {torch.cuda.is_available()}')"

永久环境变量设置（可选）

将以下内容添加到 ~/.bashrc：

复制代码

echo 'export CUDA_HOME=/usr/local/cuda' >> ~/.bashrc
echo 'export PATH=/usr/local/cuda/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc

注意事项

所有操作均在Jetson设备本地执行（非主机PC）
编译过程可能需要10-30分钟，请耐心等待
如果内存不足，可增加交换空间
确保网络连接稳定，特别是git clone操作
虚拟环境创建时使用了 --system-site-packages，继承主环境包

下一步：编译CUDA扩展

配置好PyTorch环境后，继续编译RTGS项目所需的CUDA扩展：

复制代码

# 编译Baseline CUDA扩展
cd /home/jetson/aaa-rtgs/Baseline/submodules/diff-gaussian-rasterization
pip install -e .

# 编译MonoRTGS CUDA扩展
cd /home/jetson/aaa-rtgs/MonoRTGS/submodules/diff-gaussian-rasterization
pip install -e .