### TensorRT 优化 ONNX 模型性能教程 在 Python 中使用 TensorRT 来优化 ONNX 模型的性能是一个常见的需求。以下是实现这一目标的具体方法： #### 安装必要的库 为了能够运行 TensorRT 并加载 ONNX 模型，需要安装 `pyparsing` 和 `tensorrt` 库。可以通过以下命令完成安装： ```bash pip install pycuda tensorrt onnx graphsurgeon ``` 注意：TensorRT 的版本应与 NVIDIA 驱动程序兼容[^1]。 --- #### 导入所需模块并初始化环境 在脚本中导入所需的模块，并设置 GPU 上下文以便后续操作。 ```python import tensorrt as trt from cuda import cudart import numpy as np import os ``` 确保 CUDA 已正确配置到环境中，因为 TensorRT 基于 CUDA 运行[^1]。 --- #### 加载 ONNX 模型 通过 TensorRT 提供的解析器来读取 ONNX 文件的内容。 ```python TRT_LOGGER = trt.Logger(trt.Logger.WARNING) def build_engine(onnx_file_path, engine_file_path=""): with trt.Builder(TRT_LOGGER) as builder, \ builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) as network, \ trt.OnnxParser(network, TRT_LOGGER) as parser: config = builder.create_builder_config() config.max_workspace_size = 1 << 30 # 设置最大工作空间大小为 1GB with open(onnx_file_path, 'rb') as model: if not parser.parse(model.read()): for error in range(parser.num_errors): print(parser.get_error(error)) return None engine = builder.build_serialized_network(network, config) if engine and engine_file_path != "": with open(engine_file_path, "wb") as f: f.write(engine) return engine ``` 此函数会构建一个序列化的引擎文件用于加速推理过程。 --- #### 将模型转换为可执行计划 (Engine Plan) 调用上面定义的方法生成 `.engine` 文件作为最终产物。 ```python onnx_model_path = "./model.onnx" trt_engine_path = "./model.trt" if not os.path.exists(trt_engine_path): serialized_engine = build_engine(onnx_model_path, trt_engine_path) else: with open(trt_engine_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime: serialized_engine = runtime.deserialize_cuda_engine(f.read()) ``` 如果之前已经存在对应的 Engine，则可以直接反序列化；否则重新创建它。 --- #### 执行推理流程 最后一步就是利用建立起来的 TensorRT 引擎来进行实际预测任务。 ```python class HostDeviceMem(object): def __init__(self, host_mem, device_mem): self.host = host_mem self.device = device_mem def __str__(self): return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device) def __repr__(self): return self.__str__() def allocate_buffers(engine): inputs = [] outputs = [] bindings = [] stream = cudart.cudaStreamCreate() for binding in engine: size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size dtype = trt.nptype(engine.get_binding_dtype(binding)) host_mem = cuda.pagelocked_empty(size, dtype) device_mem = cuda.mem_alloc(host_mem.nbytes) bindings.append(int(device_mem)) if engine.binding_is_input(binding): inputs.append(HostDeviceMem(host_mem, device_mem)) else: outputs.append(HostDeviceMem(host_mem, device_mem)) return inputs, outputs, bindings, stream context = serialized_engine.create_execution_context() inputs, outputs, bindings, stream = allocate_buffers(serialized_engine) input_data = np.random.rand(*[serialized_engine.max_batch_size]+list(context.get_binding_shape(0))).astype(np.float32).ravel()[:len(inputs[0].host)] np.copyto(inputs[0].host, input_data) [cudart.cudaMemcpy(inp.device, inp.host, inp.host.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice) for inp in inputs] context.execute_v2(bindings=bindings) [cudart.cudaMemcpy(out.host, out.device, out.host.nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost) for out in outputs] print(outputs[0].host.reshape(-1)[:10]) # 输出前十个结果查看效果 ``` 以上代码展示了完整的端到端解决方案，从加载预训练模型到最后获得推断输出[^1]。 --- #### 注意事项 - **精度损失**：由于 FP16 或 INT8 等量化技术的应用，在某些情况下可能会引入轻微误差。 - **硬件支持**：确认所使用的显卡具备足够的计算能力（建议至少 RTX 系列及以上）。 - **动态形状处理**：对于输入尺寸不固定的场景需额外调整网络结构参数以适应不同分辨率的要求。 ---