1. 为什么 Orin 上装 PyTorch 是件让人头疼的事如果你手里有一块 NVIDIA Orin 系列的开发板比如 Orin Nano、Orin NX 或者 AGX Orin并且已经刷好了 JetPack 6.0那你大概率已经体会过一件事在这块板子上装 PyTorch 和 Torchvision跟在普通 x86 服务器上pip install完全是两码事。我自己前后在 Orin Nano 8GB 和 AGX Orin 64GB 上都折腾过好几轮踩过的坑从“pip 装完 import 直接段错误”到“torchvision 编译到一半内存爆掉”再到“CUDA 版本对不上导致算子静默回退到 CPU”基本上能遇到的都遇到了。这篇内容就是把我这几轮折腾的经验完整梳理出来给正在或者准备在 Orin JetPack 6.0 上跑 PyTorch 2.4 Torchvision 0.19 的朋友一条相对清晰的路径。先说清楚这套组合的定位。JetPack 6.0 对应的是 L4T R36.x底层是 Ubuntu 22.04CUDA 版本是 12.2cuDNN 是 8.9TensorRT 是 8.6。PyTorch 2.4 是 2024 年下半年的版本Torchvision 0.19 是跟它配套的。这套组合在 Orin 上属于比较新的搭配NVIDIA 官方论坛上针对 JetPack 6.0 的预编译 wheel 包发布节奏比 JetPack 5.x 慢一些所以很多时候你得自己编译或者用社区维护的源。这篇文章适合三类人第一类是在 Orin 上做边缘 AI 部署、需要 PyTorch 做推理或微调的工程师第二类是在 Jetson 上做机器人、自动驾驶感知相关开发的研究人员第三类就是单纯想在这块板子上把深度学习环境跑起来的学习者。不管你是哪一类只要你的目标平台是 Orin JetPack 6.0下面的内容应该都能帮你省下不少时间。2. 环境底子摸清楚再动手2.1 先确认你的 JetPack 和 L4T 版本很多人一上来就急着装 PyTorch结果装完发现各种诡异问题根源就是没搞清楚自己板子上到底跑的是什么版本。JetPack 6.0 是一个打包概念它包含 L4T、CUDA、cuDNN、TensorRT 等一堆组件。你要做的第一件事是确认 L4T 的具体版本号。打开终端执行cat /etc/nv_tegra_release你会看到类似# R36 (release), REVISION: 3.0这样的输出。JetPack 6.0 对应的是 R36.3.0JetPack 6.0 的早期版本可能是 R36.2.0。这个版本号非常关键因为后面选 wheel 包或者编译时ABI 兼容性跟它直接相关。接着确认 CUDA 版本nvcc --versionJetPack 6.0 自带的是 CUDA 12.2。如果你看到的是 11.x那说明你刷的不是 JetPack 6.0而是 5.x那整套流程都不一样了。再确认 Python 版本python3 --versionJetPack 6.0 默认是 Python 3.10。这个也要记住因为 PyTorch wheel 包是跟 Python 版本绑定的。2.2 内存和存储的现实约束Orin 系列不同型号的内存差异很大。Orin Nano 有 4GB 和 8GB 两个版本Orin NX 有 8GB 和 16GBAGX Orin 从 32GB 到 64GB。这个内存大小直接决定了你能不能在本机编译 PyTorch。我实测下来在 Orin Nano 8GB 上编译 PyTorch 2.4 基本是不可能的编译到一半就会因为内存不足被 OOM Killer 干掉。Orin NX 16GB 可以勉强编译但需要把 swap 开到足够大而且编译时间会非常长可能要四五个小时。AGX Orin 32GB 以上编译就比较从容了。存储方面PyTorch 编译过程中产生的中间文件可能占用 20GB 以上的空间加上 CUDA 工具链本身建议至少留出 40GB 的可用空间。用df -h看一下根分区或者你打算放编译目录的分区。提示如果你用的是 Orin Nano强烈建议不要在本机编译直接用预编译的 wheel 包。省下来的时间够你做好几轮模型测试了。2.3 系统依赖的预装在动手装 PyTorch 之前有几个系统级的依赖需要先装好否则后面编译或者运行时会报各种找不到库的错误sudo apt-get update sudo apt-get install -y python3-pip python3-dev python3-venv \ build-essential cmake git libopenblas-dev libblas-dev \ liblapack-dev libjpeg-dev zlib1g-dev libpython3-dev \ libavcodec-dev libavformat-dev libswscale-dev这几个包里面libopenblas-dev和liblapack-dev是给 PyTorch 的线性代数运算用的libjpeg-dev和zlib1g-dev是 Torchvision 处理图像需要的libavcodec-dev这一组是视频解码相关的。少装一个后面就可能卡住。3. 安装路线的选择与取舍3.1 三条路官方 wheel、社区 wheel、源码编译在 Orin 上装 PyTorch本质上就三条路。第一条是 NVIDIA 官方为 Jetson 提供的预编译 wheel。NVIDIA 的开发者论坛上会定期发布针对特定 JetPack 版本的 PyTorch wheel这些 wheel 是专门为 aarch64 CUDA 编译的装上去就能用。但问题是官方 wheel 的版本更新往往滞后JetPack 6.0 刚出来那会儿官方只提供到 PyTorch 2.2 或 2.32.4 的 wheel 要等一段时间。第二条是社区维护的 wheel 源。比如一些第三方维护的 Jetson 专用 PyPI 镜像会跟进比较新的版本。用这种方式的好处是省事坏处是来源可靠性需要自己判断而且有时候 wheel 的编译选项跟你的需求不完全匹配。第三条就是从源码编译。这是最灵活的你可以自己控制编译选项比如是否启用 cuDNN、是否启用分布式训练支持等。但代价是时间长、对内存要求高、容易出错。我的建议是先看官方有没有你需要的版本的 wheel有就用官方的没有的话如果你内存够16GB 以上可以考虑源码编译如果内存不够就找社区 wheel 或者降级到官方支持的版本。3.2 为什么版本匹配这么要命PyTorch、Torchvision、CUDA、cuDNN 这四个东西的版本必须严格匹配。PyTorch 2.4 编译时链接的是 CUDA 12.2 的运行时如果你系统里的 CUDA 是 12.1 或者 12.3可能就会出现符号找不到的问题。Torchvision 0.19 又依赖特定版本的 PyTorch版本不对会直接 import 失败。更隐蔽的一个问题是 cuDNN。JetPack 6.0 自带 cuDNN 8.9PyTorch 2.4 在编译时如果检测到的 cuDNN 版本跟运行时不一致某些算子可能会静默回退到 CPU 实现你跑起来不报错但速度慢得离谱。这个问题我在 Orin NX 上遇到过排查了很久才发现是 cuDNN 版本的问题。3.3 虚拟环境还是系统环境我强烈建议用 Python 虚拟环境不要直接装在系统 Python 里。原因很简单JetPack 系统里有一些预装的 Python 包是跟系统组件绑定的你如果在系统环境里乱装东西可能会把某些系统工具搞坏。创建虚拟环境python3 -m venv ~/pytorch-env source ~/pytorch-env/bin/activate pip install --upgrade pip setuptools wheel虚拟环境建好之后后面所有的操作都在这个环境里进行。这样即使装崩了删掉虚拟环境重来就行不会影响系统。4. 实操从零到能跑通一个模型4.1 方案A使用预编译 wheel 快速安装如果你找到了对应 JetPack 6.0 Python 3.10 aarch64 的 PyTorch 2.4 wheel安装过程就很简单pip install torch-2.4.0-cp310-cp310-linux_aarch64.whl装完之后验证python3 -c import torch; print(torch.__version__); print(torch.cuda.is_available())如果输出2.4.0和True那基本就成功了。但别急着高兴还要进一步验证 CUDA 算子是否真的能用python3 -c import torch; a torch.randn(1000, 1000).cuda(); b torch.randn(1000, 1000).cuda(); c a b; print(c.sum())这个测试会实际调用 GPU 做矩阵乘法。如果这一步报错或者卡住说明 CUDA 后端有问题。Torchvision 0.19 的 wheel 也是类似pip install torchvision-0.19.0-cp310-cp310-linux_aarch64.whl验证python3 -c import torchvision; print(torchvision.__version__)4.2 方案B源码编译 PyTorch 2.4如果你必须自己编译下面是完整的流程。先确认你的内存和 swapfree -h如果物理内存小于 16GB先把 swap 加大sudo fallocate -l 16G /swapfile sudo chmod 600 /swapfile sudo mkswap /swapfile sudo swapon /swapfile然后确认 swap 生效free -h接下来拉源码。PyTorch 2.4 对应的 tag 是v2.4.0git clone --recursive --branch v2.4.0 https://github.com/pytorch/pytorch.git cd pytorch git submodule sync git submodule update --init --recursive设置编译环境变量。这一步非常关键直接影响编译出来的 PyTorch 能不能用 GPUexport USE_CUDA1 export USE_CUDNN1 export USE_NCCL1 export USE_DISTRIBUTED1 export TORCH_CUDA_ARCH_LIST8.7 export MAX_JOBS4 export USE_MKLDNN0 export USE_QNNPACK0这里解释几个关键点。TORCH_CUDA_ARCH_LIST8.7是因为 Orin 系列的 GPU 架构是 Amperecompute capability 是 8.7。如果你不设置这个PyTorch 会为所有支持的架构编译编译时间会翻好几倍。MAX_JOBS4是限制并行编译的任务数Orin 的 CPU 核心数有限设太高反而会因为内存不足崩溃。USE_MKLDNN0是因为 MKLDNN 是给 x86 用的在 ARM 上没意义关掉能省不少编译时间。开始编译python3 setup.py bdist_wheel这个过程在 Orin NX 16GB 上大概需要 3 到 5 个小时在 AGX Orin 上大概 1.5 到 2 小时。编译完成后wheel 文件在dist/目录下。安装pip install dist/torch-2.4.0*.whl4.3 编译 Torchvision 0.19Torchvision 的编译相对简单但依赖 PyTorch 已经装好。先拉源码git clone --branch v0.19.0 https://github.com/pytorch/vision.git cd vision设置环境变量export FORCE_CUDA1 export TORCH_CUDA_ARCH_LIST8.7 export MAX_JOBS4编译安装python3 setup.py bdist_wheel pip install dist/torchvision-0.19.0*.whlTorchvision 编译时间比 PyTorch 短很多大概 20 到 40 分钟。4.4 验证整套环境装完之后跑一个完整的验证脚本import torch import torchvision import torchvision.transforms as T from PIL import Image import numpy as np print(PyTorch:, torch.__version__) print(Torchvision:, torchvision.__version__) print(CUDA available:, torch.cuda.is_available()) print(CUDA version:, torch.version.cuda) print(cuDNN version:, torch.backends.cudnn.version()) print(Device name:, torch.cuda.get_device_name(0)) # 测试张量运算 x torch.randn(3, 3).cuda() y torch.randn(3, 3).cuda() z torch.mm(x, y) print(Matrix multiply on GPU:, z.device) # 测试 Torchvision 的模型 from torchvision.models import resnet18 model resnet18(pretrainedFalse).cuda().eval() dummy torch.randn(1, 3, 224, 224).cuda() with torch.no_grad(): out model(dummy) print(ResNet18 output shape:, out.shape) # 测试图像变换 transform T.Compose([ T.Resize(256), T.CenterCrop(224), T.ToTensor(), T.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) img Image.fromarray(np.random.randint(0, 255, (300, 300, 3), dtypenp.uint8)) tensor transform(img) print(Transformed tensor shape:, tensor.shape)这个脚本覆盖了 PyTorch 的基本张量运算、CUDA 可用性、Torchvision 的模型加载和图像变换。如果全部通过说明环境基本没问题。5. 那些让人抓狂的坑和排查方法5.1 import torch 报段错误这是最常见的问题之一。表现是import torch直接 Segmentation fault连报错信息都没有。原因通常是 PyTorch 编译时链接的库跟系统里的库版本不匹配。排查方法用ldd看 PyTorch 的共享库依赖ldd $(python3 -c import torch; print(torch.__file__) | sed s/__init__.py/libtorch.so/)重点看libcudart.so、libcudnn.so、libcublas.so这几个的路径。如果指向的是系统里旧版本的库而不是 JetPack 6.0 自带的那就是问题所在。解决办法是确保LD_LIBRARY_PATH包含 JetPack 的 CUDA 库路径export LD_LIBRARY_PATH/usr/local/cuda/lib64:$LD_LIBRARY_PATH5.2 CUDA 可用但算子报错有时候torch.cuda.is_available()返回 True但一跑实际算子就报CUDA error: no kernel image is available for execution on the device。这是因为编译时的TORCH_CUDA_ARCH_LIST没有包含 8.7。解决办法是重新编译确保TORCH_CUDA_ARCH_LIST8.7。如果你用的是预编译 wheel那说明这个 wheel 不是为 Orin 编译的换一个。5.3 Torchvision 找不到匹配的 PyTorchTorchvision 在 import 时会检查 PyTorch 版本。如果版本不匹配会报RuntimeError: PyTorch version mismatch。这个问题的根源是 Torchvision 0.19 要求 PyTorch 2.4.x如果你装的是 2.3 或者 2.5就会报错。解决办法就是确保两者版本严格对应。Torchvision 0.19 对应 PyTorch 2.4Torchvision 0.18 对应 PyTorch 2.3以此类推。5.4 编译过程中内存不足前面提到过Orin Nano 8GB 编译 PyTorch 基本不可能。如果你在编译过程中看到g: fatal error: Killed signal terminated program cc1plus那就是 OOM 了。除了加大 swap还可以减少并行编译任务数export MAX_JOBS2或者只编译你需要的部分比如关掉分布式支持export USE_DISTRIBUTED05.5 常见问题速查表问题现象可能原因排查方法解决方式import torch 段错误库版本不匹配ldd 检查依赖设置 LD_LIBRARY_PATHCUDA 算子报 no kernel image编译架构不对检查 TORCH_CUDA_ARCH_LIST重新编译设为 8.7Torchvision 版本不匹配PyTorch 版本不对检查两者版本号安装对应版本编译时 OOM内存不足free -h 查看加大 swap减少 MAX_JOBSGPU 算子静默回退 CPUcuDNN 版本不一致检查 torch.backends.cudnn.version()确保 cuDNN 版本匹配pip 安装 wheel 报平台不兼容wheel 不是 aarch64检查 wheel 文件名下载正确的 wheel6. 性能调优与日常使用建议6.1 让 Orin 的 GPU 跑满装好 PyTorch 只是第一步要让 Orin 的 GPU 真正发挥性能还需要做一些设置。首先是电源模式Orin 默认可能是低功耗模式GPU 频率被限制sudo nvpmodel -m 0 sudo jetson_clocksnvpmodel -m 0是切换到最大性能模式jetson_clocks是锁定最高频率。这两个命令在跑推理或者训练之前执行能明显提升速度。但要注意散热Orin 在高负载下发热不小如果散热不好会降频。6.2 混合精度训练Orin 的 GPU 支持 FP16 和 BF16用混合精度能显著提升训练速度、降低显存占用from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for data, target in dataloader: data, target data.cuda(), target.cuda() optimizer.zero_grad() with autocast(dtypetorch.float16): output model(data) loss criterion(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()在 Orin 上FP16 的加速比大概在 1.5 到 2 倍左右具体取决于模型结构。6.3 显存管理Orin 的显存是和系统内存共享的所以显存管理跟 x86 独显不太一样。你可以通过torch.cuda.memory_summary()查看显存使用情况。如果遇到显存不足可以尝试torch.cuda.empty_cache()或者调整 batch size。在 Orin Nano 8GB 上ResNet50 的 batch size 大概只能到 16 左右再大就 OOM 了。6.4 模型导出与 TensorRT 加速如果你只是做推理PyTorch 原生推理在 Orin 上并不是最快的。可以考虑导出成 ONNX 或者直接用 TensorRTtorch.onnx.export(model, dummy_input, model.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch}, output: {0: batch}})导出成 ONNX 之后再用 TensorRT 做进一步优化推理速度通常能再提升 2 到 3 倍。不过这是另一个话题了这里不展开。7. 我踩过的几个印象深刻的坑第一个坑是在 Orin Nano 上编译 PyTorch编译到 80% 的时候 OOM整个编译过程白费。后来我学乖了编译之前先free -h确认内存不够就加 swap。而且编译的时候不要开其他占内存的程序浏览器都不行。第二个坑是 cuDNN 版本问题。我装完 PyTorch 之后跑一个卷积网络速度比预期慢很多但没有任何报错。后来用torch.backends.cudnn.version()一看发现 PyTorch 链接的是系统里一个旧版本的 cuDNN而不是 JetPack 6.0 自带的。解决办法是在虚拟环境的激活脚本里显式设置LD_LIBRARY_PATH。第三个坑是 Torchvision 的编译。Torchvision 0.19 在编译时需要 PyTorch 的源码路径如果你只装了 PyTorch 的 wheel 而没有源码编译会失败。解决办法是设置TORCHVISION_INCLUDE和TORCHVISION_LIBRARY环境变量指向 PyTorch 的安装路径。第四个坑是 Python 版本。JetPack 6.0 默认是 Python 3.10但如果你自己装了其他版本的 Pythonpip 可能会装到错误的 Python 环境里。一定要用python3 -m pip而不是直接pip确保装到正确的环境。最后分享一个小技巧如果你在 Orin 上做开发建议把常用的环境变量写进~/.bashrc或者虚拟环境的activate脚本里这样每次打开终端就不用重新设置。特别是LD_LIBRARY_PATH和TORCH_CUDA_ARCH_LIST这两个忘了设置就会出各种奇怪的问题。这套环境配好之后我在 Orin 上跑 YOLOv8 的推理大概能到 30 到 40 FPSFP16640x640 输入跑 ResNet50 的分类大概 100 FPS 左右。对于边缘设备来说这个性能已经相当可用了。后续如果要做模型量化或者 TensorRT 部署那又是另一套流程有机会再单独聊。