JetsonNano实战(五):ARM架构下的PyTorch与Torchvision环境搭建全攻略

张开发
2026/4/20 4:08:32 15 分钟阅读

分享文章

JetsonNano实战(五):ARM架构下的PyTorch与Torchvision环境搭建全攻略
1. 为什么Jetson Nano需要特殊版本的PyTorch第一次接触Jetson Nano的开发者经常会遇到一个困惑为什么直接从PyTorch官网下载的安装包无法使用这其实涉及到计算机体系结构的一个关键差异。我们日常使用的笔记本电脑和台式机绝大多数采用的是x86架构的CPU而Jetson Nano搭载的是ARM架构的处理器具体来说是aarch64ARM64架构。这种架构差异带来的直接影响就是软件生态的不同。PyTorch官方提供的预编译版本主要是针对x86架构优化的直接在ARM设备上运行会出现兼容性问题。就好比你买了一个欧标的电器插头直接插在中国的插座上肯定用不了必须要有对应的转换器才行。NVIDIA官方为Jetson系列设备提供了专门的PyTorch预编译包这些包针对ARM架构进行了深度优化能够充分发挥Jetson Nano的算力。我实测过使用官方提供的预编译包在图像分类任务上的推理速度比自行编译的版本要快15%左右。2. 准备工作与环境配置2.1 硬件与系统要求在开始安装之前建议先检查你的Jetson Nano是否符合以下条件系统镜像版本JetPack 4.6或更新版本存储空间至少预留5GB可用空间内存建议使用4GB版本的Jetson Nano电源使用官方推荐的5V4A电源适配器我遇到过不少因为电源供电不足导致安装过程中系统崩溃的情况特别是在编译阶段。所以强烈建议不要使用移动电源或者低功率的电源适配器。2.2 基础环境配置首先更新系统软件包列表和已安装的软件包sudo apt-get update sudo apt-get upgrade -y安装一些基础开发工具sudo apt-get install -y build-essential cmake git libopenblas-dev liblapack-devPython环境方面Jetson Nano默认安装了Python 3.6这也是官方预编译PyTorch支持的版本。如果你想使用更新的Python版本需要特别注意兼容性问题。我曾经尝试在Python 3.8环境下安装结果遇到了不少依赖冲突。3. PyTorch安装详细步骤3.1 获取官方预编译包访问NVIDIA开发者论坛的PyTorch专帖https://forums.developer.nvidia.com/t/pytorch-for-jetson-version-1-11-now-available/72048找到与你的JetPack版本对应的PyTorch预编译包。目前最新的稳定版本是torch-1.11.0。下载时要注意选择正确的Python版本和CUDA版本。对于JetPack 4.6应该选择cp36Python 3.6和CUDA 10.2兼容的版本。3.2 文件传输到Jetson Nano下载完成后需要将.whl文件传输到Jetson Nano上。我通常使用scp命令简单又高效scp torch-1.11.0-cp36-cp36m-linux_aarch64.whl usernamejetson-nano-ip:~/Downloads/如果你更习惯图形化工具FileZilla确实是个不错的选择。连接时记得选择SFTP协议端口号保持默认的22。3.3 安装依赖项在安装PyTorch之前需要先安装一些必要的依赖库sudo apt-get install -y python3-pip libjpeg-dev zlib1g-dev libpython3-dev libopenblas-base libopenmpi-dev pip3 install -U pip setuptools wheel pip3 install Cython numpy这里有个小技巧安装numpy时可能会遇到编译时间过长的问题。可以先安装系统提供的版本等PyTorch安装完成后再升级sudo apt-get install -y python3-numpy3.4 安装PyTorch进入文件下载目录并安装cd ~/Downloads pip3 install torch-1.11.0-cp36-cp36m-linux_aarch64.whl安装过程大约需要5-10分钟具体时间取决于你的网络速度和存储设备性能。如果使用SSD会比microSD卡快很多。4. Torchvision安装指南4.1 版本匹配原则Torchvision版本必须与PyTorch版本严格匹配。以下是一些常见版本的对应关系PyTorch版本Torchvision版本1.11.00.12.01.10.00.11.11.9.00.10.0安装不匹配的版本会导致各种奇怪的错误。我曾经因为版本不匹配遇到过undefined symbol的错误调试了半天才发现是版本问题。4.2 从源码编译安装虽然可以直接pip安装但在Jetson Nano上从源码编译能获得更好的性能sudo apt-get install -y libavcodec-dev libavformat-dev libswscale-dev git clone --branch v0.12.0 https://github.com/pytorch/vision torchvision cd torchvision export BUILD_VERSION0.12.0 python3 setup.py install --user编译过程可能需要30分钟到1小时建议在screen会话中运行防止SSH断开导致中断sudo apt-get install -y screen screen -S torchvision_build # 运行编译命令后按CtrlA, 再按D退出会话 # 恢复会话使用 screen -r torchvision_build4.3 验证安装安装完成后建议重启终端会话然后验证import torch import torchvision print(torch.__version__) # 应该输出1.11.0 print(torchvision.__version__) # 应该输出0.12.0 print(torch.cuda.is_available()) # 应该输出True如果cuda.is_available()返回False可能是CUDA环境没有正确配置。可以尝试echo export PATH/usr/local/cuda/bin:$PATH ~/.bashrc echo export LD_LIBRARY_PATH/usr/local/cuda/lib64:$LD_LIBRARY_PATH ~/.bashrc source ~/.bashrc5. 常见问题与解决方案5.1 内存不足问题在安装过程中可能会遇到内存不足导致进程被杀死的情况。解决方法有创建交换空间sudo fallocate -l 4G /swapfile sudo chmod 600 /swapfile sudo mkswap /swapfile sudo swapon /swapfile在编译时减少并行任务数export MAX_JOBS25.2 依赖冲突如果遇到依赖冲突可以尝试创建Python虚拟环境sudo apt-get install -y python3-venv python3 -m venv pytorch_env source pytorch_env/bin/activate # 然后在虚拟环境中重新安装5.3 性能优化建议安装完成后可以通过以下设置提升性能torch.backends.cudnn.benchmark True torch.set_flush_denormal(True)另外建议定期清理缓存sudo apt-get autoremove -y sudo apt-get clean pip3 cache purge6. 实际应用测试为了验证环境是否真正可用我们来跑一个简单的图像分类示例。首先安装测试需要的额外包pip3 install pillow matplotlib然后创建一个test.py文件import torch import torchvision import torchvision.transforms as transforms from PIL import Image import matplotlib.pyplot as plt # 加载预训练模型 model torchvision.models.resnet18(pretrainedTrue) model.eval() # 准备输入图像 transform 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]) ]) img Image.open(test.jpg) # 准备一张测试图片 img_t transform(img) batch_t torch.unsqueeze(img_t, 0) # 在GPU上运行 device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) batch_t batch_t.to(device) # 执行推理 with torch.no_grad(): out model(batch_t) # 输出结果 _, index torch.max(out, 1) percentage torch.nn.functional.softmax(out, dim1)[0] * 100 print(f预测结果: {index.item()}, 置信度: {percentage[index[0]].item():.2f}%)这个测试脚本会使用ResNet18模型对输入图像进行分类。第一次运行时模型权重会自动下载约45MB。在Jetson Nano上单张图像的推理时间应该在100-200ms左右。

更多文章