BlueXIII

Technical note

从PyTorch到ONNX和TensorRT的YOLO部署

阅读 本文由AI辅助编写
AI

在NVIDIAGPU上部署YOLO时，最终采用了TensorRTEngine加Ultralytics推理。作出选择之前，需要先弄清.pt、.onnx和.engine分别解决什么问题，以及ONNXRuntime和TensorRT对环境的要求。

参考

三种模型文件

.pt是PyTorch检查点，也是训练阶段最方便保留的格式。使用Ultralytics加载时，模型结构、权重和任务相关信息都由原来的Python环境处理，适合继续训练、验证和直接推理。

.onnx保存的是标准化计算图，不依赖原来的PyTorch执行过程。它可以交给ONNXRuntime的CPU、CUDA或TensorRTExecutionProvider，也方便接入其他支持ONNX的推理引擎。

这里的.engine专指TensorRT序列化引擎，也叫plan。它已经针对TensorRT版本、平台和GPU进行编译，启动时不再重新解析完整的ONNX图，但兼容范围比ONNX窄。默认生成的engine不能随意换到另一种GPU、操作系统或TensorRT版本上运行，部署环境变化后通常需要重新构建。

选择TensorRTEngine加Ultralytics推理，主要是为了保持调用代码一致。这个选择不代表engine在所有环境都更快；脱离目标硬件、输入尺寸、批量大小和精度谈固定的性能百分比没有意义，延迟和吞吐量还是要在实际部署环境中比较。

导出ONNX

先在训练环境安装Ultralytics和ONNX依赖：

python -m pip install --upgrade ultralytics onnx

使用Python导出：

from ultralytics import YOLO

model = YOLO("best.pt")
model.export(
    format="onnx",
    imgsz=640,
    dynamic=False,
    simplify=True,
)

输入尺寸固定时保持dynamic=False，部署配置更简单；确实需要接受多种尺寸时再启用动态形状。opset不指定时由Ultralytics选择当前支持的版本，只有目标运行时有明确限制时才固定。

导出后先检查模型格式：

import onnx

model = onnx.load("best.onnx")
onnx.checker.check_model(model)

使用ONNXRuntime

CPU和GPU包二选一安装，不要同时留在同一个环境中：

# CPU
python -m pip install onnxruntime

# NVIDIAGPU
python -m pip install onnxruntime-gpu

PyPI上的onnxruntime-gpu默认对应CUDA12.x。实际安装前还要根据ONNXRuntime的CUDAExecutionProvider兼容表核对CUDA和cuDNN主版本，不能只看nvidia-smi显示的CUDA上限。

确认当前环境真正加载了哪些ExecutionProvider：

import onnxruntime as ort

print(ort.__version__)
print(ort.get_available_providers())

创建CUDA推理会话时把CPU放在后面作为回退：

import onnxruntime as ort

session = ort.InferenceSession(
    "best.onnx",
    providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
)
print(session.get_providers())

get_available_providers()只能说明当前安装包包含哪些Provider，session.get_providers()用来确认这个Session接受了哪些Provider。两者都不能单独证明每个节点实际分配到了GPU，还要留意启动日志中是否有CUDA或cuDNN动态库加载失败。

TensorRTExecutionProvider

ONNXRuntime也可以把支持的子图交给TensorRT，不支持的节点继续交给CUDA。Provider顺序要把TensorRT放在CUDA前面：

import onnxruntime as ort

providers = [
    (
        "TensorrtExecutionProvider",
        {
            "device_id": 0,
            "trt_fp16_enable": True,
            "trt_engine_cache_enable": True,
            "trt_engine_cache_path": "./trt-cache",
        },
    ),
    ("CUDAExecutionProvider", {"device_id": 0}),
    "CPUExecutionProvider",
]

session = ort.InferenceSession("best.onnx", providers=providers)
print(session.get_providers())

第一次创建Session会编译TensorRT子图，冷启动可能很慢。启用engine缓存后，后续启动可以直接复用；模型、ONNXRuntime、TensorRT、构建参数或GPU发生变化时应清理旧缓存重新生成。

导出TensorRTEngine

Ultralytics可以从.pt直接生成engine：

from ultralytics import YOLO

model = YOLO("best.pt")
model.export(
    format="engine",
    imgsz=640,
    quantize=16,
    workspace=4,
    verbose=True,
)

当前Ultralytics文档使用quantize=16表示FP16。旧版本常见的half=True仍可能可用，但已经属于兼容参数。workspace=4表示最多使用4GiB工作空间；显存不紧张时也可以省略，让TensorRT自动决定。

engine会针对构建时的GPU和TensorRT环境选择内核，最好直接在最终部署机器上导出。若需要跨机器分发，先根据TensorRT的版本兼容和硬件兼容选项确认边界，不能把普通engine当成ONNX一样复制使用。

安装TensorRT运行时

TensorRT提供PyPI、Debian/RPM、Tar包和容器等安装方式。Python项目只使用API时，PyPI最省事；需要trtexec、C++头文件或同时保留多个系统版本时，再选择系统包或Tar包。

例如现有环境是CUDA12.x，可以明确安装对应的Python包：

python -m pip install --upgrade tensorrt-cu12

验证Python接口和CUDA初始化：

import tensorrt as trt

print(trt.__version__)
assert trt.Builder(trt.Logger())

PyPI安装不包含trtexec。如果使用Debian、RPM、Tar包或容器安装，再用下面的命令检查CLI：

trtexec --help

TensorRT、CUDA、驱动、Python和GPU的支持范围变化很快，安装时直接从支持矩阵选择一组匹配版本，不要把某一组TensorRT、CUDA和Python版本当作长期安装规则。

自己编写CUDA扩展时，CUDA_HOME应指向Toolkit根目录，而不是它的bin目录：

export CUDA_HOME=/usr/local/cuda-X.Y

上面的Ultralytics、ONNXRuntime和TensorRTPython路径没有直接使用PyCUDA，不需要为了运行engine额外安装。只有自己的推理代码明确import pycuda时，再按该项目的依赖补上。

使用Ultralytics统一调用

如果不想分别维护ONNXRuntime和TensorRT推理代码，可以继续用Ultralytics加载三种文件。模型路径变化，预测接口基本一致：

from ultralytics import YOLO

pt_model = YOLO("best.pt")
onnx_model = YOLO("best.onnx")
engine_model = YOLO("best.engine")

pt_model.predict(source="image.jpg", device=0)
onnx_model.predict(source="image.jpg", device=0)
engine_model.predict(source="image.jpg", device=0)

这种方式部署简单，但运行环境仍要安装Ultralytics及其依赖。追求更小的运行环境或需要C++集成时，直接使用ONNXRuntime或TensorRT原生API更合适。