手把手教你用Conda在Jetson AGX Orin上配置PyTorch 1.12和Torchvision 0.16.0

在Jetson AGX Orin上使用Conda配置PyTorch 1.12和Torchvision 0.16.0的完整指南

Jetson AGX Orin作为NVIDIA推出的边缘计算设备，凭借其强大的AI算力，已经成为众多开发者在计算机视觉和深度学习领域的首选硬件平台。然而，由于ARM架构的特殊性，许多开发者在使用过程中遇到了PyTorch环境配置的挑战。本文将详细介绍如何通过Conda这一强大的环境管理工具，在Jetson AGX Orin上搭建PyTorch 1.12和Torchvision 0.16.0的开发环境，帮助开发者避开常见的"坑"，快速进入项目开发阶段。

1. 环境准备与基础配置

在开始安装之前，我们需要确保Jetson AGX Orin的系统环境已经准备就绪。首先检查系统版本和CUDA驱动是否正常安装：

uname -a cat /etc/nv_tegra_release nvidia-smi

这些命令将分别显示系统架构信息、JetPack版本以及CUDA驱动状态。Jetson AGX Orin通常预装了JetPack SDK，其中包含了CUDA、cuDNN等必要的深度学习库。

接下来，我们需要安装Conda。由于Jetson采用ARM架构，我们推荐使用Miniforge的ARM版本：

wget https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-aarch64.sh bash Miniforge3-Linux-aarch64.sh

安装完成后，初始化Conda环境并创建一个新的虚拟环境：

conda create -n pytorch_1.12 python=3.8 -y conda activate pytorch_1.12

为什么选择Python 3.8？NVIDIA官方提供的PyTorch预编译包对Python版本有特定要求，3.8是目前最稳定的兼容版本。

2. CUDA工具包与PyTorch安装

在Jetson平台上安装PyTorch需要特别注意CUDA版本的选择。首先查看可用的CUDA工具包版本：

conda search cudatoolkit

对于PyTorch 1.12，我们推荐使用CUDA 11.4版本：

conda install cudatoolkit=11.4 -c conda-forge

安装完成后，验证CUDA是否正常工作：

nvcc --version

接下来，我们需要从NVIDIA官方下载适用于Jetson AGX Orin的PyTorch预编译包。访问NVIDIA PyTorch下载页面找到对应版本。

下载完成后，使用pip安装：

pip install torch-1.12.0a0+nnnnnnnnnnnn_linux_aarch64.whl

安装完成后，验证PyTorch是否正确识别了CUDA：

import torch print(torch.__version__) print(torch.cuda.is_available())

3. Torchvision源码编译与安装

由于Torchvision没有官方预编译的ARM版本，我们需要从源码编译安装。首先确保系统已经安装了必要的编译工具：

sudo apt-get install libjpeg-dev zlib1g-dev libpython3-dev libavcodec-dev libavformat-dev libswscale-dev

然后克隆Torchvision源码库，注意选择与PyTorch 1.12兼容的v0.16.0版本：

git clone --branch v0.16.0 https://github.com/pytorch/vision torchvision cd torchvision

设置环境变量并开始编译：

export BUILD_VERSION=v0.16.0 python3 setup.py install --user

编译过程可能需要较长时间（约30-60分钟），具体取决于Jetson AGX Orin的性能配置。编译完成后，验证安装：

import torchvision print(torchvision.__version__)

4. 常见问题排查与性能优化

在实际安装过程中，开发者可能会遇到各种问题。以下是一些常见问题及其解决方案：

问题1：编译Torchvision时报错

error: command 'aarch64-linux-gnu-gcc' failed with exit status 1

解决方案：确保已安装所有必要的开发依赖：

sudo apt-get install build-essential python3-dev

问题2：PyTorch无法识别CUDA

torch.cuda.is_available()返回False

解决方案：检查CUDA版本是否匹配，并确保正确安装了NVIDIA驱动。

性能优化建议：

1. 在编译Torchvision时，可以设置以下环境变量加速编译：

export MAX_JOBS=$(nproc)

2. 使用swap空间避免内存不足：

sudo fallocate -l 4G /swapfile sudo chmod 600 /swapfile sudo mkswap /swapfile sudo swapon /swapfile

3. 在运行深度学习模型时，设置以下环境变量优化性能：

export OMP_NUM_THREADS=1 export OPENBLAS_CORETYPE=ARMV8

5. 实际应用测试与验证

为了验证我们的安装是否成功，我们可以运行一个简单的图像分类示例。首先安装测试所需的额外依赖：

pip install pillow matplotlib

然后创建一个测试脚本test.py：

import torch import torchvision from torchvision import transforms from PIL import Image print("PyTorch版本:", torch.__version__) print("Torchvision版本:", torchvision.__version__) print("CUDA可用:", torch.cuda.is_available()) # 加载预训练模型 model = torchvision.models.resnet18(pretrained=True) model.eval() # 准备输入图像 input_image = Image.open("test.jpg") preprocess = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) input_tensor = preprocess(input_image) input_batch = input_tensor.unsqueeze(0) # 将模型和输入数据移动到GPU if torch.cuda.is_available(): model = model.cuda() input_batch = input_batch.cuda() # 执行推理 with torch.no_grad(): output = model(input_batch) # 输出结果 print("推理完成，输出形状:", output.shape)

运行此脚本将验证PyTorch和Torchvision是否正常工作，并展示如何在Jetson AGX Orin上运行深度学习模型。