4. NVIDIA TensorRT加速座舱视觉模型推理
任务说明
在座舱AI场景中，驾驶员监控（DMS）与乘客状态识别等视觉模型需要满足实时性要求。本案例将完成一个最小可复现的TensorRT加速闭环：将PyTorch预训练的ResNet-18分类模型转换为ONNX，再通过TensorRT构建FP16精度的推理引擎，最终对比TensorRT与PyTorch的推理延迟，验证加速效果。
环境依赖
- Ubuntu 20.04+ / JetPack 5.0+（车规级Orin平台）
- TensorRT 8.5+（含trtexec命令行工具）
- PyTorch 2.0+（仅用于导出ONNX与基准对比）
- CUDA 11.8+、cuDNN 8.6+
- Python 3.8+，需安装pycuda、onnx、onnxruntime（可选）
实现步骤
整体流程分为三步：导出ONNX → 构建TensorRT引擎 → 执行推理并验证。以下代码在一个Python脚本中完成全部闭环，输入为随机生成的模拟座舱图像张量（3×224×224），输出为分类置信度向量。
关键完整代码
import torch
import torchvision.models as models
import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import time
# ---------- 1. 导出ONNX ----------
model = models.resnet18(pretrained=False)
model.eval()
dummy_input = torch.randn(1, 3, 224, 224)
onnx_path = "cabin_model.onnx"
torch.onnx.export(model, dummy_input, onnx_path,
input_names=["input"], output_names=["output"],
dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}},
opset_version=17)
print("ONNX导出完成")
# ---------- 2. 构建TensorRT FP16引擎 ----------
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, TRT_LOGGER)
with open(onnx_path, "rb") as f:
if not parser.parse(f.read()):
for i in range(parser.num_errors):
print(parser.get_error(i))
raise RuntimeError("ONNX解析失败")
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30) # 1GB
config.set_flag(trt.BuilderFlag.FP16)
engine_path = "cabin_model_fp16.engine"
serialized_engine = builder.build_serialized_network(network, config)
with open(engine_path, "wb") as f:
f.write(serialized_engine)
print("TensorRT引擎构建完成")
# ---------- 3. 加载引擎并推理 ----------
runtime = trt.Runtime(TRT_LOGGER)
engine = runtime.deserialize_cuda_engine(serialized_engine)
context = engine.create_execution_context()
# 分配显存
input_shape = (1, 3, 224, 224)
output_shape = (1, 1000)
input_size = trt.volume(input_shape) * np.float32().itemsize
output_size = trt.volume(output_shape) * np.float32().itemsize
d_input = cuda.mem_alloc(input_size)
d_output = cuda.mem_alloc(output_size)
# 模拟座舱输入图像（随机噪声）
input_data = np.random.randn(*input_shape).astype(np.float32)
output_data = np.empty(output_shape, dtype=np.float32)
# 绑定TensorRT上下文
context.set_input_shape("input", input_shape)
bindings = [d_input, d_output]
# 预热
cuda.memcpy_htod(d_input, input_data)
context.execute_v2(bindings)
cuda.memcpy_dtoh(output_data, d_output)
cuda.Context.synchronize()
# 正式推理计时（100次取平均）
n_iter = 100
start = time.time()
for _ in range(n_iter):
cuda.memcpy_htod(d_input, input_data)
context.execute_v2(bindings)
cuda.memcpy_dtoh(output_data, d_output)
cuda.Context.synchronize()
trt_time = (time.time() - start) / n_iter * 1000 # ms
# ---------- 4. PyTorch CPU/GPU基准对比 ----------
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
torch_input = torch.from_numpy(input_data).to(device)
# 预热
with torch.no_grad():
for _ in range(10):
_ = model(torch_input)
# 计时
start = time.time()
with torch.no_grad():
for _ in range(n_iter):
_ = model(torch_input)
torch.cuda.synchronize() if device == "cuda" else None
torch_time = (time.time() - start) / n_iter * 1000
print(f"TensorRT FP16 平均延迟: {trt_time:.2f} ms")
print(f"PyTorch ({device}) 平均延迟: {torch_time:.2f} ms")
print(f"加速比: {torch_time / trt_time:.2f}x")
print(f"输出shape: {output_data.shape}, 前3个值: {output_data[:3]}")运行方式
# 安装依赖
pip install torch torchvision onnx pycuda tensorrt
# 执行脚本（需在NVIDIA GPU/Orin设备上）
python trt_cabin_demo.py验证方法与通过标准
- 引擎构建成功：脚本输出“TensorRT引擎构建完成”，且当前目录生成cabin_model_fp16.engine文件（约30MB）。
- 推理输出正确：输出shape为(1, 1000)，数值为有限浮点数（非NaN/Inf）。
- 加速效果达标：TensorRT FP16延迟显著低于PyTorch GPU模式（通常快2～5倍）。若在无GPU环境运行，PyTorch将回退到CPU，此时加速比会更高，但仍需确认TensorRT推理正常完成。
- 确定性预期输出：由于输入为随机噪声，输出值无法预先确定，但可验证输出张量每个元素均为有限值，且argmax索引在0～999之间。
常见问题
- 显存不足：将workspace从1GB降至256MB（
1 << 28），或减小batch size。 - FP16精度损失：若分类结果与FP32差异过大，可保留FP32引擎对比，或使用INT8量化（需校准数据集）。
- 动态shape报错：本案例已设置动态batch，若需固定shape可移除dynamic_axes参数。