### 使用TensorRT运行ONNX模型的Python代码示例 为了使用TensorRT加速ONNX模型推理，在Python环境中需先安装必要的库，包括`onnxruntime-gpu`和`tensorrt`。这些工具能够显著提升深度学习模型在GPU上的执行效率[^1]。 #### 安装依赖项 确保已安装支持CUDA的NVIDIA驱动程序以及cuDNN库之后，可以通过pip安装所需包： ```bash pip install nvidia-pyindex pip install --upgrade onnx-graphsurgeon pip install --extra-index-url https://pypi.ngc.nvidia.com nvidia-tensorrt pip install onnxruntime_tensorrt ``` #### 将ONNX模型转换为TensorRT引擎 下面是一个简单的脚本用于加载ONNX文件并将其编译成适用于TensorRT的优化版本： ```python import tensorrt as trt from polygraphy.backend.trt import CreateConfig, EngineFromNetwork, NetworkFromOnnxPath, SaveEngine import os def build_engine(onnx_file_path, engine_file_path): TRT_LOGGER = trt.Logger(trt.Logger.WARNING) network_creation_flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) builder = trt.Builder(TRT_LOGGER) explicit_batch_network_flags = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) network = builder.create_network(explicit_batch_network_flags ) config = builder.create_builder_config() parser = trt.OnnxParser(network, TRT_LOGGER) with open(onnx_file_path, 'rb') as model: if not parser.parse(model.read()): print('Failed to parse the ONNX file.') for error in range(parser.num_errors): print(parser.get_error(error)) return None config.max_workspace_size = 1 << 30 # 1GB engine = builder.build_serialized_network(network, config) with open(engine_file_path, "wb") as f: f.write(bytearray(engine)) build_engine("model.onnx", "model.engine") ``` 此函数接受两个参数：一个是输入的`.onnx`路径；另一个是要保存生成的TensorRT序列化网络（即engine）的目标位置。 #### 运行推断过程 一旦拥有了TensorRT引擎文件，则可以在实际应用中快速高效地完成预测任务： ```python import numpy as np import pycuda.driver as cuda import pycuda.autoinit import tensorrt as trt 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 = cuda.Stream() for binding in engine: size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size dtype = trt.nptype(engine.get_binding_dtype(binding)) # Allocate host and device buffers host_mem = cuda.pagelocked_empty(size, dtype) device_mem = cuda.mem_alloc(host_mem.nbytes) # Append the device buffer to device bindings. bindings.append(int(device_mem)) # Append to the appropriate list. 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 with open("model.engine", "rb") as f, trt.Runtime(TRT_LOGGER) as runtime: engine = runtime.deserialize_cuda_engine(f.read()) context = engine.create_execution_context() inputs, outputs, bindings, stream = allocate_buffers(engine) # 假设我们有一个形状为(batch_size,)的数据集作为输入数据 input_data = np.random.rand(*engine.get_binding_shape(0)).astype(np.float32) np.copyto(inputs[0].host, input_data.ravel()) [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs] context.execute_async_v2(bindings=bindings, stream_handle=stream.handle) [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs] stream.synchronize() output = [out.host.reshape(engine.get_binding_shape(i+1)) for i,out in enumerate(outputs)] print(output) ``` 这段代码展示了如何创建上下文对象来管理资源分配，并通过异步调用来提高性能。最后打印输出结果以供进一步处理或分析[^2]。