3. 模型优化
任务目标
将训练好的车辆检测模型（ONNX格式）通过TensorRT进行优化加速，在嵌入式GPU平台（如Jetson Orin）上获得可部署的TensorRT引擎，并验证推理加速效果。
环境依赖
- 硬件：NVIDIA Jetson Orin / Xavier 系列，或任意支持TensorRT的NVIDIA GPU
- 软件：TensorRT 8.5+、CUDA 11.4+、cuDNN 8.5+
- Python：3.8+，需安装
tensorrt、pycuda、onnx、numpy、Pillow - 输入模型：车辆检测ONNX模型（输入为
[1,3,640,640]的RGB图像，输出为检测结果张量）
实现方案
本案例采用TensorRT的ONNX解析器将车辆检测模型转换为TensorRT引擎，核心优化策略包括：FP16精度量化、动态batch支持、以及通过工作空间限制控制显存占用。转换后的引擎将保存为本地文件供后续部署直接加载，避免重复转换开销。
关键完整代码
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import time
from PIL import Image
# ---------- 1. TensorRT 引擎构建 ----------
class VehicleDetectorTRT:
def __init__(self, onnx_path, engine_path, precision='fp16', max_batch=4):
self.onnx_path = onnx_path
self.engine_path = engine_path
self.precision = precision
self.max_batch = max_batch
self.logger = trt.Logger(trt.Logger.WARNING)
self.engine = None
self.context = None
self.stream = None
self.inputs = []
self.outputs = []
self.bindings = []
self._load_or_build_engine()
def _build_engine(self):
"""从ONNX构建TensorRT引擎"""
builder = trt.Builder(self.logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, self.logger)
with open(self.onnx_path, 'rb') as f:
if not parser.parse(f.read()):
for i in range(parser.num_errors):
print(f"ONNX解析错误: {parser.get_error(i)}")
raise RuntimeError("ONNX模型解析失败")
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30) # 1GB工作空间
if self.precision == 'fp16' and builder.platform_has_fast_fp16:
config.set_flag(trt.BuilderFlag.FP16)
print("已启用FP16精度")
# 动态batch配置
profile = builder.create_optimization_profile()
input_tensor = network.get_input(0)
input_name = input_tensor.name
profile.set_shape(input_name,
min=(1, 3, 640, 640),
opt=(self.max_batch, 3, 640, 640),
max=(self.max_batch, 3, 640, 640))
config.add_optimization_profile(profile)
print("开始构建TensorRT引擎（可能需要数分钟）...")
self.engine = builder.build_serialized_network(network, config)
if self.engine is None:
raise RuntimeError("引擎构建失败")
with open(self.engine_path, 'wb') as f:
f.write(self.engine)
print(f"引擎已保存至: {self.engine_path}")
def _load_or_build_engine(self):
"""加载已有引擎或重新构建"""
import os
if os.path.exists(self.engine_path):
print(f"加载已有引擎: {self.engine_path}")
with open(self.engine_path, 'rb') as f:
self.engine = self.logger.deserialize(f.read())
else:
self._build_engine()
with open(self.engine_path, 'rb') as f:
self.engine = self.logger.deserialize(f.read())
self.context = self.engine.create_execution_context()
self.stream = cuda.Stream()
self._setup_bindings()
def _setup_bindings(self):
"""分配输入输出缓冲区"""
for i in range(self.engine.num_io_tensors):
tensor_name = self.engine.get_tensor_name(i)
tensor_shape = self.engine.get_tensor_shape(tensor_name)
tensor_dtype = trt.nptype(self.engine.get_tensor_dtype(tensor_name))
# 动态维度处理
if -1 in tensor_shape:
tensor_shape = (self.max_batch, 3, 640, 640) if i == 0 else tensor_shape
self.context.set_input_shape(tensor_name, tensor_shape)
size = int(np.prod(tensor_shape))
host_mem = cuda.pagelocked_empty(size, dtype=tensor_dtype)
device_mem = cuda.mem_alloc(host_mem.nbytes)
self.bindings.append(device_mem)
if self.engine.get_tensor_mode(tensor_name) == trt.TensorIOMode.INPUT:
self.inputs.append({'name': tensor_name, 'host': host_mem, 'device': device_mem, 'shape': tensor_shape})
else:
self.outputs.append({'name': tensor_name, 'host': host_mem, 'device': device_mem, 'shape': tensor_shape})
def preprocess(self, image_path):
"""图像预处理：resize到640x640并归一化"""
img = Image.open(image_path).convert('RGB')
img = img.resize((640, 640), Image.BILINEAR)
img_array = np.array(img, dtype=np.float32) / 255.0
img_array = img_array.transpose(2, 0, 1) # HWC -> CHW
img_array = np.expand_dims(img_array, axis=0) # 添加batch维度
return np.ascontiguousarray(img_array)
def infer(self, image_path):
"""执行推理"""
input_data = self.preprocess(image_path)
input_host = self.inputs[0]['host']
input_host[:] = input_data.reshape(-1)
# 异步传输与推理
cuda.memcpy_htod_async(self.inputs[0]['device'], input_host, self.stream)
self.context.execute_async_v2(bindings=self.bindings, stream_handle=self.stream.handle)
cuda.memcpy_dtoh_async(self.outputs[0]['host'], self.outputs[0]['device'], self.stream)
self.stream.synchronize()
output_shape = self.outputs[0]['shape']
output_data = self.outputs[0]['host'].reshape(output_shape)
return output_data
def benchmark(self, image_path, iterations=100):
"""性能基准测试"""
# 预热
for _ in range(10):
self.infer(image_path)
# 正式计时
start = time.time()
for _ in range(iterations):
self.infer(image_path)
elapsed = time.time() - start
avg_latency = elapsed / iterations * 1000 # 转换为毫秒
fps = iterations / elapsed
print(f"平均延迟: {avg_latency:.2f} ms")
print(f"吞吐量: {fps:.1f} FPS")
return avg_latency, fps
# ---------- 2. 使用示例 ----------
if __name__ == "__main__":
ONNX_PATH = "vehicle_detector.onnx" # 输入ONNX模型路径
ENGINE_PATH = "vehicle_detector.trt" # 输出TensorRT引擎路径
# 构建并加载引擎
detector = VehicleDetectorTRT(ONNX_PATH, ENGINE_PATH, precision='fp16')
# 单张图片推理
test_image = "test_car.jpg"
output = detector.infer(test_image)
print(f"推理输出shape: {output.shape}")
# 性能测试
detector.benchmark(test_image, iterations=50)
运行方式
- 将车辆检测ONNX模型重命名为
vehicle_detector.onnx并放置于脚本同目录 - 准备一张包含车辆的测试图片
test_car.jpg - 执行脚本：
python trt_vehicle_detector.py
验证方法与通过标准
- 引擎构建成功：控制台输出“引擎已保存至: vehicle_detector.trt”，且生成的文件大小在数十MB级别（取决于原模型大小）
- 推理输出正确：输出张量的shape与ONNX模型输出定义一致，数值范围合理（检测框坐标、置信度等）
- 性能达标：在Jetson Orin上，FP16精度下车辆检测模型的推理延迟应低于原始ONNX Runtime推理延迟的50%（例如原30ms降至15ms以下），FPS提升超过2倍
- 精度损失可接受：FP16与FP32的检测结果（IoU > 0.5的检测框数量）差异不超过5%
常见问题
- 显存不足：降低
max_batch或减小工作空间限制（set_memory_pool_limit参数） - ONNX解析失败：确认ONNX模型使用opset ≥ 11，且不包含TensorRT不支持的算子（如动态shape的某些实现）
- 首次构建耗时较长：FP16引擎构建可能需要5-10分钟，属正常现象；后续加载引擎文件只需数秒