4、TensorRT/ONNX Runtime部署：模型转换、动态batch、多流并发优化

模型训练完了，精度也达标了。但说实话，这只是第一步。在车载座舱里，芯片资源就那么点，摄像头数据还一帧一帧地来。怎么把训练好的模型塞进芯片，让它跑得快、跑得稳？这就是部署要干的事。

我个人习惯把部署分成两派：TensorRT 和 ONNX Runtime。前者是英伟达的看家本领，专治 GPU 推理；后者是微软牵头的开放标准，CPU、GPU、NPU 通吃。车载座舱里，这两样都得会。

核心观点：部署不是简单的模型格式转换，而是对计算图的重构和优化。你想想看，训练时我们关心精度，部署时我们关心延迟和吞吐。这是两套不同的优化目标。

4.1 模型转换：从训练框架到推理引擎

模型转换，说白了就是把 PyTorch 或 TensorFlow 的模型，变成 TensorRT 或 ONNX 能吃的格式。我刚开始做的时候，以为就是调个 API 完事。结果第一次转 ONNX 就报错，查了半天，原来是 PyTorch 里有个自定义算子，ONNX 不认识。

这里我总结了一套标准流程：

  1. 导出 ONNX：从 PyTorch 用 torch.onnx.export，从 TensorFlow 用 tf2onnx。
  2. ONNX 检查：用 onnx.checker 验证模型结构，用 onnxruntime 跑一遍推理，确保精度对齐。
  3. TensorRT 转换：用 trtexec 或 Python API 把 ONNX 转成 TensorRT 引擎。
  4. 精度验证：对比转换前后的输出，误差通常在 1% 以内。
# PyTorch 转 ONNX 示例
import torch
import torch.onnx

model = torch.load('model.pth')
model.eval()

dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model,
    dummy_input,
    'model.onnx',
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    },
    opset_version=11
)
print("ONNX 导出完成")
我的经验：导出 ONNX 时，opset_version 别用太新。车载芯片的驱动版本往往滞后，用 opset 11 或 12 最稳。我曾经用 opset 15 导出，结果在 Jetson Xavier 上跑不了，折腾了两天。

4.2 动态 batch：应对不确定的输入数量

车载座舱的场景很特殊。有时候只检测一个人，有时候后排坐满了一家子。如果模型只支持固定 batch size，要么浪费算力，要么装不下。动态 batch 就是解决这个问题的。

动态 batch 的核心思想：推理引擎在运行时根据实际输入数量，动态调整计算图。TensorRT 和 ONNX Runtime 都支持，但实现方式不同。

特性 TensorRT ONNX Runtime
配置方式 构建引擎时指定优化范围 运行时设置输入形状
性能 预编译，延迟低 动态编译，首次调用慢
灵活性 需预设 batch 范围 完全动态，无限制
内存管理 预分配，内存可控 动态分配，可能碎片化

在 TensorRT 里配置动态 batch，需要在构建引擎时指定优化 profile：

import tensorrt as trt

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
profile = builder.create_optimization_profile()

# 定义动态 batch 范围：最小 1，最优 4，最大 8
profile.set_shape('input', (1, 3, 224, 224), (4, 3, 224, 224), (8, 3, 224, 224))
config = builder.create_builder_config()
config.add_optimization_profile(profile)

# 序列化引擎
engine = builder.build_serialized_network(network, config)
注意：动态 batch 的优化范围不要设得太宽。比如你设成 1 到 32，TensorRT 会尝试优化所有中间尺寸，导致引擎文件巨大，编译时间超长。我一般根据实际场景设：座舱里最多 6 个人，我就设 1 到 8，留点余量。

4.3 多流并发优化：榨干芯片的每一分算力

车载座舱里，多个算法同时跑是常态。人脸检测、手势识别、疲劳监测……每个模型都要占资源。如果串行跑，延迟叠加，用户体验极差。多流并发就是让这些模型同时跑，互不干扰。

多流并发的核心是 CUDA Stream。每个 stream 是一个独立的执行队列，GPU 可以在多个 stream 之间切换，实现真正的并行。

我画了一张图，帮你理解多流并发的架构：

多流并发推理架构 输入数据流（摄像头帧、麦克风音频） Stream 分配器（负载均衡） Stream 1 人脸检测模型 batch=4 Stream 2 手势识别模型 batch=2 Stream 3 疲劳监测模型 batch=1 GPU 并行执行（CUDA 核心 + Tensor Core） Stream 1 卷积 → Stream 2 全连接 → Stream 3 池化 → 交替执行 每个 Stream 独立管理上下文，互不阻塞。GPU 在 Stream 间自动切换，最大化利用率。

实现多流并发的代码其实不复杂。核心是创建多个 CUDA stream，然后每个模型绑定一个 stream：

import tensorrt as trt
import pycuda.driver as cuda

# 创建多个 CUDA stream
streams = [cuda.Stream() for _ in range(3)]

# 为每个模型创建独立的执行上下文
contexts = []
for i in range(3):
    engine = load_engine(f'model_{i}.trt')
    context = engine.create_execution_context()
    context.set_optimization_profile_async(0, streams[i].handle)
    contexts.append(context)

# 多流并发推理
def infer_multi_stream(inputs):
    for i, (ctx, stream) in enumerate(zip(contexts, streams)):
        # 异步拷贝输入
        cuda.memcpy_htod_async(d_inputs[i], inputs[i], stream)
        # 异步执行推理
        ctx.execute_async_v2(bindings[i], stream.handle, None)
        # 异步拷贝输出
        cuda.memcpy_dtoh_async(outputs[i], d_outputs[i], stream)
    
    # 同步所有 stream
    cuda.Context.synchronize()
    return outputs
性能调优技巧：多流并发不是越多越好。GPU 的并行能力有限，stream 太多反而增加调度开销。我一般根据 GPU 的 SM 数量来定：Jetson Orin 有 8 个 SM，我就开 4 到 6 个 stream。你可以用 nvidia-smi 看 GPU 利用率，调到 90% 以上就差不多了。

4.4 实战避坑指南

部署过程中坑不少。我踩过的，都记下来了：

  • 精度对齐问题：转换后精度掉 1-2% 是正常的。如果掉太多，检查算子兼容性。我遇到过 PyTorch 的 nn.Upsample 转 ONNX 后行为不一致，换成 torch.nn.functional.interpolate 才解决。
  • 动态 batch 内存爆炸：如果 batch 上限设得太大，TensorRT 会预分配大量显存。我建议用 profile.set_shape 时，上限不超过实际需求的 2 倍。
  • 多流并发死锁：多个 stream 同时访问同一个模型上下文，会导致死锁。每个 stream 必须有自己的 context，不能共享。
  • ONNX 算子缺失：车载芯片的 ONNX Runtime 版本往往较老，不支持新算子。我一般用 onnxruntime.transformers 做优化，或者手动替换算子。
曾经踩过的坑：有一次我在 Jetson Xavier 上部署多流并发，结果 GPU 利用率只有 30%。查了两天，发现是 CPU 端的数据预处理成了瓶颈。摄像头帧的 JPEG 解码占用了大量 CPU 时间，导致 GPU 一直在等数据。解决方案是用 GPU 的 nvjpeg 做硬件解码，把预处理也扔到 GPU 上。利用率直接飙到 85%。

部署这件事，说白了就是和硬件打交道。每个芯片都有自己的脾气，TensorRT 和 ONNX Runtime 只是工具。真正值钱的，是你对硬件特性的理解，以及踩坑后积累的经验。嗯，今天就先聊到这，这些内容够你消化一阵子了。

← 上一章 📖 返回目录 下一章 →