实验室里有几台 NVIDIA JETSON ORIN以及最新购入的 NVIDIA THOR设备，本文记录一下这些推理设备的基本信息和使用方法。

1. Jetson 平台简介

Jetson 是 NVIDIA 推出的边缘计算平台。其本质是 SoC(System on Chip，片上系统)，集成了 CPU、GPU、内存等多个组件，专为边缘计算和人工智能推理设计。其设计理念是将需要在服务器/显卡上运行的 AI 推理，放到低功耗、实时、嵌入式的场景中。

1.1 特点

其CPU、GPU、NPU、内存都集成在一个同一芯片上；内存是统一内存，CPU/GPU/NPU全部共享同一块内存；系统采用ARM架构（aarch64），与传统 x86 架构不同。

1.2 软件生态的不同

Jetson 完全不同于传统的 PC 计算平台，并非“装UBUNTU，装驱动，装CUDA”这种流程。
而是使用 JetPack SDK 来进行系统和软件的安装和配置。一切内容直接跟刷机的 JetPack 版本绑定。

JetPack = Linux 系统 + GPU 驱动 + CUDA + cuDNN + TensorRT

2. 常见的 Jetson 设备

2.1 JETSON NANO

Jetson NANO 是 NVIDIA 于 2019 年发布的一款入门级边缘计算设备。其定位是低功耗、低成本的 AI 推理平台，适合教育、开发和轻量级应用。功耗仅 5W-10W.

规格参数:

  • GPU：Maxwell 架构
  • AI算力：472 GFLOPS（0.5 TFLOPS）
  • 内存：4GB LPDDR4
  • CPU：四核 ARM Cortex-A57

说实话，这个性能确实比较低了，跑个YOLO模型FPS都不高。

2.2 JETSON ORIN NX

2.3 NVIDIA THOR

NVIDIA THOR 是 NVIDIA 于 2024 年发布的一款高性能边缘计算设备。差不多3W一台，性能可见一斑。不过整个机身非常大，最开始看到我以为是个PC上的显卡嘞。

规格参数:

  • GPU：Blackwell 架构
  • AI算力：2070 TFLOPS(FP4)
  • 显存带宽：273 GB/s
  • CPU：14核

3. 网络计算流程

3.1 网络计算的本质

无论是CNN还是Transformer，推理的核心都是大量的线性代数+少量的逐元素操作。

  • 卷积/全连接: 矩阵乘法(GEMM)或者类似的张量乘加
  • BatchNorm: 逐元素的线性变换
  • Relu等激活函数: 逐元素非线性变换
  • Pooling/Resize/NMS: 一些规则性/几何性操作。（NMS有时候是在CPU上做的）

GPU在做的就是把输入张量和权重张量按某种规则做乘加运算，再做一些逐元素操作，最后输出结果张量。

3.2 普通平台（PC）上的计算流

3.2.1 用户编写层（使用pytorch/tensorflow等框架）

假设我们用pytorch写了一个简单的前向过程

1
x -> conv1 -> bn1 -> relu -> conv2 -> ...

这段代码本质是在”搭积木+调用“。做的工作是：组织要使用哪些算子，决定它们的执行顺序，调用最后实现的底层算子。

3.2.2 Pytorch层

Pytorch主要做下面几件事：

  • Tensor 抽象与管理：dtype、shape、device等
  • 算子调度：根据Tensor的device，选择在CPU/GPU上执行对应的算子实现（eg. CUDA kernel）
  • 内存管理：分配GPU显存、管理Tensor生命周期
  • 训练相关：autograd自动求导、优化器、反向传播等

3.2.3 CUDA层

更细节的说，pytorch调用某个算子时，cuDNN提供最优实现（选择不同算法），由CUDA kernel在GPU上执行。

3.3 Jetson 平台上的计算流

Jetson平台上，整体计算流和PC上类似，但有一些区别：

1
pytorch层  →  onnx层  →  TensorRT层  →  CUDA层

3.3.1 ONNX 层

Jetson平台上，通常会先把pytorch模型导出为ONNX格式。
ONNX是一个开放的神经网络交换格式，目的是实现不同深度学习框架之间的互操作性。这是一方面的理由，对于边缘处理设备来说，ONNX把动态的pytorch模型转换为静态图，可以更好地进行优化和部署。
这个转换包括：

  • 网络拓扑固定
  • 算子顺序固定
  • 权重常量化
  • 尽量消除python控制流影响

这些都是为了之后的TensorRT优化做准备（全局优化）。

3.3.2 TensorRT 层

TensorRT像是一个编译器，将网络编译成一份针对特定GPU的执行计划。
其中的关键工作包括：

  • 算子融合：例如 Conv + BN + Relu 融合为一个算子，减少内存读写/少启动kernel
  • 精度优化：用更低精度（FP16/INT8）计算，提高吞吐
  • 内存规划：直接全局上来分析中间结果的生命周期，减少内存占用
  • 算法选择和调度固化：选择最优算法，并把调度顺序固化，减少运行时开销

3.3.3 开始推理

不同于PC上每次运行会重新调度实现，Jetson上TensorRT生成的引擎文件已经把所有细节都固化好了，形成了一个Engine，只需要传入输入张量，执行算子管线即可。
所以，相比于PC上的灵活性，Jetson平台上更注重性能和效率。其针对的问题包括：

  • kernel启动开销
  • 内存搬运开销
  • CPU调度太多
  • 延迟抖动