1. 环境准备与基础配置

拿到Jetson Orin NX的第一件事，就是做好基础环境配置。我建议直接从NVIDIA官网下载最新的JetPack SDK，这是最省事的方案。我实测过JetPack 5.1.2版本，稳定性相当不错，CUDA 11.4、cuDNN 8.6这些核心组件都预装好了，省去很多折腾时间。

安装完系统后，第一件事就是更新软件源。Jetson Orin NX用的是ARM架构，很多软件包需要从特定源获取。我习惯用清华源，速度会快很多：

sudo cp /etc/apt/sources.list /etc/apt/sources.list.backup
sudo sed -i 's/ports.ubuntu.com/mirrors.tuna.tsinghua.edu.cn/' /etc/apt/sources.list
sudo apt update

接下来安装jtop这个神器，它能实时监控Jetson的运行状态，包括CPU、GPU、内存使用情况，还有JetPack各组件的版本信息：

sudo apt install python3-pip
sudo -H pip3 install -U jetson-stats

装完后直接运行jtop，就能看到漂亮的监控界面。这里要特别注意CUDA版本是否正常显示，我遇到过因为驱动问题导致CUDA不识别的情况，这时候需要重新安装NVIDIA驱动。

2. PyTorch GPU版本安装

Jetson平台安装PyTorch有个坑：直接用pip install会装成CPU版本。必须从NVIDIA官网下载预编译的whl文件。我建议去PyTorch for Jetson页面找对应版本，JetPack 5.1.2对应的是PyTorch v2.1.0。

下载whl文件后，安装命令很简单：

pip install torch-2.1.0a0+41361538.nv23.06-cp38-cp38-linux_aarch64.whl

这里有个细节要注意：如果下载速度慢，可以加上清华源加速：-i https://pypi.tuna.tsinghua.edu.cn/simple

安装完成后一定要验证GPU是否可用：

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

如果最后一行显示"Orin"，说明PyTorch已经正确识别到你的Jetson Orin NX的GPU了。我遇到过torch.cuda.is_available()返回False的情况，一般都是因为CUDA版本不匹配，需要重新检查JetPack版本。

3. Torchvision编译安装

PyTorch装好了，接下来要安装torchvision。这里要注意版本匹配：PyTorch 2.1.0需要torchvision 0.16.x。官方没有提供预编译包，必须从源码编译。

我试过直接clone源码，但github经常抽风，下载到一半就断连。更稳妥的方法是直接去torchvision的github仓库，下载v0.16.2的zip包，然后本地解压。

编译前需要先安装一堆依赖：

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

然后设置环境变量并编译：

cd torchvision
export BUILD_VERSION=0.16.2
sudo python3 setup.py install

编译过程大概需要10-20分钟，取决于你的网络速度。完成后用pip list检查，应该能看到torchvision 0.16.2。

我在这里踩过一个坑：编译时内存不足导致失败。Jetson Orin NX的内存虽然不小，但编译时还是建议关闭其他程序，或者临时增加swap空间：

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

4. 环境验证与问题排查

环境装好后一定要全面测试。我写了个简单的测试脚本，检查所有关键组件：

import torch
import torchvision
import numpy as np

print("PyTorch版本:", torch.__version__)
print("Torchvision版本:", torchvision.__version__)
print("CUDA可用:", torch.cuda.is_available())
print("GPU名称:", torch.cuda.get_device_name(0))

# 测试GPU计算
x = torch.randn(3, 3).cuda()
y = torch.randn(3, 3).cuda()
z = x + y
print("GPU计算测试通过")

# 测试卷积操作
conv = torch.nn.Conv2d(3, 16, 3).cuda()
input = torch.randn(1, 3, 32, 32).cuda()
output = conv(input)
print("卷积计算测试通过")

如果所有测试都通过，说明PyTorch环境配置成功。常见的几个问题：一是内存不足，二是依赖库缺失，三是版本不匹配。

内存不足的问题可以通过增加swap解决，依赖库缺失一般会报很明显的错误信息，按照提示安装即可。版本不匹配最麻烦，需要仔细检查PyTorch、torchvision、CUDA三者的版本对应关系。

5. TensorRT环境配置

TensorRT是NVIDIA的推理加速库，Jetson Orin NX上已经预装了。可以用以下命令检查版本：

dpkg -l | grep tensorrt

如果显示8.x版本，说明已经安装。但为了充分发挥性能，我建议更新到最新版本：

sudo apt update
sudo apt install tensorrt

TensorRT和PyTorch的配合使用需要一些技巧。首先要把PyTorch模型转换成ONNX格式，然后再用TensorRT优化。我写了个转换脚本：

import torch
import torchvision
import tensorrt as trt

# 示例模型
model = torchvision.models.resnet18(pretrained=True)
model.eval()

# 导出ONNX
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(model, dummy_input, "resnet18.onnx", verbose=True)

然后用trtexec工具转换成TensorRT引擎：

trtexec --onnx=resnet18.onnx --saveEngine=resnet18.engine

这个过程中可能会遇到算子不支持的问题，需要根据错误信息调整模型结构或者使用自定义算子。

6. YOLOv5部署实战

现在我们来实战部署YOLOv5模型。首先克隆官方仓库：

git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txt

安装依赖时要注意：有些包需要特定版本，比如numpy最好用1.23.5，新版本可能会有兼容性问题。

下载预训练权重：

wget https://github.com/ultralytics/yolov5/releases/download/v6.0/yolov5s.pt

测试推理：

import torch
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
results = model('https://ultralytics.com/images/bus.jpg')
results.print()

如果能看到检测结果，说明YOLOv5运行正常。接下来用TensorRT加速：

python export.py --weights yolov5s.pt --include engine --device 0

导出过程中可能会遇到问题，最常见的是ONNX版本兼容性。建议使用ONNX 1.12.0版本：

pip install onnx==1.12.0 onnxsim==0.4.17

导出的engine文件可以直接用TensorRT加载，速度会比原始PyTorch模型快很多。

7. 性能优化技巧

Jetson Orin NX的性能潜力很大，但需要一些调优技巧。首先启用GPU频率最大化：

sudo nvpmodel -m 0
sudo jetson_clocks

这两条命令会让GPU运行在最高频率，但功耗和发热也会增加。如果需要平衡性能和功耗，可以用nvpmodel -m 1。

内存管理也很重要。Jetson Orin NX的共享内存架构意味着CPU和GPU共用内存，需要合理配置：

import torch
torch.cuda.set_per_process_memory_fraction(0.8)  # 限制GPU内存使用

推理时的batch size对性能影响很大。一般来说，batch size越大，吞吐量越高，但延迟也会增加。需要根据实际需求权衡：

# 测试不同batch size的性能
for batch_size in [1, 2, 4, 8]:
    input = torch.randn(batch_size, 3, 640, 640).cuda()
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    
    start.record()
    model(input)
    end.record()
    torch.cuda.synchronize()
    
    print(f"Batch size {batch_size}: {start.elapsed_time(end):.2f}ms")

另外，混合精度推理可以显著提升速度：

from torch.cuda.amp import autocast

with autocast():
    output = model(input)

8. 实际应用中的坑与解决方案

在实际项目中，我遇到过不少问题。其中一个常见问题是内存泄漏，特别是在长时间运行的服务中。解决方法是用torch.cuda.empty_cache()定期清理缓存：

import gc
import torch

def clean_memory():
    gc.collect()
    torch.cuda.empty_cache()

另一个问题是模型加载慢。Jetson Orin NX的存储速度相对较慢，大型模型加载需要较长时间。可以考虑将模型预先加载到内存中，或者使用更轻量的模型。

温度控制也很重要。长时间高负载运行会导致过热降频，建议监控温度：

watch -n 1 cat /sys/class/thermal/thermal_zone*/temp

如果温度过高，可以考虑加装散热片或者风扇。

最后提醒一点：Jetson Orin NX的电源管理比较敏感，建议使用官方推荐的电源适配器。电压不稳可能导致各种奇怪的问题，比如突然重启或者性能下降。