如何解决YOLO-NAS部署难题？ONNX转换与TensorRT加速完整指南

YOLO-NAS 的部署链路和普通 YOLO 模型有一个明显差异：训练仓库 super-gradients 导出的 ONNX 并不总是为 TensorRT 准备好。它可能带有动态维度、默认 NMS，或者包含一些 TensorRT 当前版本不直接支持的算子。实际落地时，比较稳妥的顺序是先导出、再清洗、最后构建 engine。本文围绕这几个步骤给出一套可复用的流程，重点解决 ONNX 转换和 TensorRT 加速过程中的高频问题。

一、ONNX 导出：先把入口参数固定下来

super-gradients 的版本不同，导出 API 和默认行为也会变化。导出前需要先确认使用的版本，4.x 和 5.x 在模型加载、导出函数上的差异会直接影响后续转换结果。比较稳定的做法是调用 prep_model_for_conversion 方法，它会完成 RepVGG 分支融合、重参数化等工作，避免把训练态结构带进 ONNX。如果不调用这个方法直接导出，计算图中可能残留大量分支，TensorRT 构建 engine 时会明显变慢，甚至触发算子不支持错误。

输入形状建议在导出阶段就固定为 [1, 3, 640, 640]。不要为了后续动态 batch 而让 ONNX 保持动态维度，动态 batch 更合适在 TensorRT optimization profile 中处理。导出时可以指定 opset_version 为 16 或 17，较高的 opset 对 LayerNorm、GELU 等算子表达更自然，但也要看目标 TensorRT 版本的支持范围。下面是一个基础的导出示例：

pip install super-gradients==4.9.0

from super_gradients.training import models

model = models.get("yolo_nas_s", pretrained_weights="coco")
model.eval()
model.cuda()
model.prep_model_for_conversion(input_size=[1, 3, 640, 640])

import torch
torch.onnx.export(
    model,
    torch.randn(1, 3, 640, 640).cuda(),
    "yolo_nas_s_raw.onnx",
    input_names=["images"],
    output_names=["output0"],
    opset_version=16,
    do_constant_folding=True,
)

执行完导出后，建议先用 ONNX Runtime 加载一次，确认输出数量和名称。YOLO-NAS 如果保留了默认后处理，输出可能已经包含 boxes、scores、class_ids 三个张量；如果关闭后处理，则输出是 decode 之前的回归与分类结果。TensorRT 对带 NMS 的 ONNX 支持并不稳定，很多时候需要在下游自己实现 decode 和后处理，因此第一步就要根据业务需求决定保留哪些输出。

二、ONNX 清洗：能折叠的折叠，该切掉的切掉

从 PyTorch 导出的 ONNX 通常会带有多余的 Shape、Gather、Concat 操作，这些节点在 TensorRT 构建时会拖慢解析速度，也可能导致部分算子无法识别。清洗计算图的第一步是用 onnxsim 做常量折叠和 identity 消除。安装 onnx-simplifier 后，指定输入形状执行简化，输出通常比原始 ONNX 更小、更干净。简化命令如下：

python -m onnxsim yolo_nas_s_raw.onnx yolo_nas_s_sim.onnx --input-shape images:1,3,640,640

简化完成后，还需要扫描一次 ONNX 中的算子类型，尤其关注 TensorRT 长期不支持的节点。可以写一段 Python 脚本遍历计算图，把 ScatterND、NonMaxSuppression、NonZero、GatherND 等节点打印出来。如果发现这些节点，优先考虑从导出侧去掉对应操作，而不是等 TensorRT 报错后再猜原因。检查脚本如下：

import onnx

model = onnx.load("yolo_nas_s_sim.onnx")
unsupported_ops = ["ScatterND", "NonMaxSuppression", "NonZero", "GatherND"]
for node in model.graph.node:
    if node.op_type in unsupported_ops:
        print("unsupported op:", node.name, node.op_type)

如果确认不需要模型内部的后处理，可以用 polygraphy surgeon 做计算图裁剪。比如只保留 decode 前的输出，把 NMS 相关分支移除，能显著提高 engine 构建成功率。polygraphy 的 prune 命令可以在不破坏输入输出的前提下删掉多余节点，也可以显式指定保留哪些输出。对于 YOLO-NAS 这种默认带完整后处理的模型，切掉后期 NMS 通常比试图转换 NMS 插件更省事。

另一个常见做法是让业务端自己写 decode。模型只输出回归偏移和分类概率，之后在 C++ 或 Python 后端完成 anchor decode、置信度过滤和 NMS。这样做虽然多了一些业务代码，但部署兼容性更好，也方便针对不同硬件调整后处理性能。TensorRT 自带的 EfficientNMS plugin 也有使用场景，但它对 ONNX 的输出格式有固定要求，YOLO-NAS 原生输出往往不能直接对接，需要额外改图适配。

三、构建 TensorRT engine：trtexec 和 Python API

最快速验证 GPU 加速效果的方式是用 TensorRT 自带的 trtexec 工具。它会自动解析 ONNX、构建 engine，并可以指定 FP16、动态 shape 等参数。命令中 minShapes、optShapes、maxShapes 分别对应最小、常用和最大输入形状，三者必须满足维度一致。对于单 batch 部署，可以直接都设为 1x3x640x640；如果需要动态 batch，就把 maxShapes 改大。一个典型的 FP16 构建命令如下：

trtexec --onnx=yolo_nas_s_sim.onnx \
    --saveEngine=yolo_nas_s_fp16.engine \
    --fp16 \
    --minShapes=images:1x3x640x640 \
    --optShapes=images:1x3x640x640 \
    --maxShapes=images:8x3x640x640

Windows 环境下如果 trtexec 没有加入 PATH，可以进入类似 C:\TensorRT-8.6.1.6\bin 的目录执行命令。Linux 服务器一般安装到 /usr/src/tensorrt 或 /opt/TensorRT，路径要按实际版本调整。trtexec 构建完成后会在终端打印 engine 的 latency 和 throughput，可以作为精度和速度的基线，不过它不会做业务后处理，只能评估模型本身的推理耗时。

如果需要在服务代码中动态构建 engine，可以使用 TensorRT Python API。相比 trtexec，Python API 能更细粒度地控制 builder、network、config 和 optimization profile。下面是一个完整构建示例，注意其中显式设置了 EXPLICIT_BATCH，并启用了 FP16：

import tensorrt as trt

logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, logger)

with open("yolo_nas_s_sim.onnx", "rb") as f:
    if not parser.parse(f.read()):
        for i in range(parser.num_errors):
            print(parser.get_error(i))
        raise RuntimeError("ONNX parse failed")

config = builder.create_builder_config()
config.set_flag(trt.BuilderFlag.FP16)
config.max_workspace_size = 1 << 30

profile = builder.create_optimization_profile()
profile.set_shape("images", (1, 3, 640, 640), (1, 3, 640, 640), (8, 3, 640, 640))
config.add_optimization_profile(profile)

engine = builder.build_serialized_network(network, config)
with open("yolo_nas_s_fp16.engine", "wb") as f:
    f.write(engine)

engine 构建完成后，加载和推理还需要处理显存分配。TensorRT 的 execute_v2 使用 bindings 列表传递输入输出指针，输入输出 buffer 的大小必须和 profile 中定义的 shape 一致。推理阶段的 Python 示例可以作为模板：

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np

engine = trt.Runtime(trt.Logger(trt.Logger.WARNING)).deserialize_cuda_engine(
    open("yolo_nas_s_fp16.engine", "rb").read()
)
context = engine.create_execution_context()
context.set_binding_shape(0, (1, 3, 640, 640))

input_buf = np.ascontiguousarray(np.random.randn(1, 3, 640, 640).astype(np.float32))
output = np.empty((1, 84, 8400), dtype=np.float32)

d_input = cuda.mem_alloc(input_buf.nbytes)
d_output = cuda.mem_alloc(output.nbytes)
bindings = [int(d_input), int(d_output)]

cuda.memcpy_htod(d_input, input_buf)
context.execute_v2(bindings)
cuda.memcpy_dtoh(output, d_output)
print(output.shape)

实际使用时，输出形状不能写死，最好通过 engine.get_binding_shape 或 context.get_binding_shape 动态获取。输入数据进入模型前还要做与训练一致的预处理，包括 letterbox 缩放、颜色通道顺序和归一化。例如训练阶段如果用的是 RGB 0 到 1 输入，部署端也必须保持相同方式，否则检测框位置和置信度会整体偏移。

四、常见报错和排查路径

ScatterND 报错是 YOLO-NAS 转 TensorRT 时最典型的问题之一。这个算子通常来自模型内部的索引或注意力结构，当前部分 TensorRT 版本无法原生支持。遇到这类错误，第一反应不应该是降级模型，而是检查 ONNX 导出链路是否可以绕开。简化计算图后仍然存在时，可以升级 TensorRT 到 8.6 或更高版本，或者通过裁剪网络去掉相关分支。若只在后处理中出现，直接在导出时关闭后处理即可解决。

engine 加载阶段出现 deserialization failed 或 plugin version mismatch，基本可以判定是 TensorRT 版本不一致。构建 engine 的 TensorRT 版本必须与线上加载 engine 的版本完全匹配，甚至小版本不同也可能导致反序列化失败。建议把 TensorRT、CUDA、驱动版本固定在同一套容器镜像中，避免构建环境和运行环境漂移。

如果使用 INT8 量化，还容易出现校准数据不足导致的精度大幅下降。YOLO-NAS 对 INT8 校准比较敏感，一般需要准备几百张真实业务场景的图片，而不是随机噪声或训练集的一个小角落。校准数据分布越接近线上输入，量化误差越小。没有充分校准数据时，更建议使用 FP16，性能和精度之间的平衡更好。

推理时 CUDA out of memory 通常由过大的 max batch 或 workspace 设置引起。构建 engine 时 maxShapes 定义得越大，GPU 显存占用越高，同时可能降低 kernel 选择空间。线上服务如果只跑单 batch，就尽量把 profile 的 max batch 限制到 2 或 4，不要为了以后可能用到的动态 batch 预留过大空间。输出 buffer 的分配也要根据实际 binding size 计算，不能按固定数值开大数组，否则容易在长时间运行时出现内存碎片和显存浪费。

YOLO-NAS 部署难的核心不是模型本身，而是从训练框架到推理引擎的中间链路。把 ONNX 导出参数固定、做一轮计算图清洗、再用 TensorRT 构建 engine，基本可以覆盖大多数落地场景。遇到不支持的算子时，优先选择让模型输出 decode 前的张量，把后处理移到业务侧，通常能获得更高的部署稳定性和更好的硬件适配性。

YOLO-NAS部署ONNX转换TensorRT加速修改时间：2026-10-04 12:28:35