### Python 推理 Engine 的实现代码示例 以下是基于 TensorRT 和 OpenVINO 的两种常见 Python 推理引擎实现方式： #### 基于 TensorRT 的推理实现 TensorRT 是一种高性能深度学习推理库，支持优化神经网络模型并将其部署到生产环境中。以下是一个简单的 TensorRT 推理实现代码示例[^3]。 ```python import numpy as np import pycuda.autoinit import pycuda.driver as cuda import tensorrt as trt class HostDeviceMem(object): """Simple helper data class that's a little nicer to use than a 2-tuple.""" 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): """Allocates all host/device buffers required for an 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 input or output 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 def do_inference(context, bindings, inputs, outputs, stream, batch_size=1): """Performs inference on a given context with provided inputs.""" # Transfer input data to the GPU. [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs] # Run inference. context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle) # Transfer predictions back from the GPU. [cuda.memcopy_dtoh_async(out.host, out.device, stream) for out in outputs] # Synchronize the stream. stream.synchronize() # Return only the host outputs. return [out.host for out in outputs] if __name__ == "__main__": TRT_LOGGER = trt.Logger(trt.Logger.WARNING) runtime = trt.Runtime(TRT_LOGGER) with open("model_data/yolov4.engine", "rb") as f: engine = runtime.deserialize_cuda_engine(f.read()) context = engine.create_execution_context() inputs, outputs, bindings, stream = allocate_buffers(engine) # Example input (replace this with actual preprocessed image data). dummy_input = np.random.rand(*inputs[0].host.shape).astype(np.float32) np.copyto(inputs[0].host, dummy_input.ravel()) result = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)[0] print(result) ``` 此代码展示了如何加载预训练的 `.engine` 文件，并通过分配内存缓冲区执行推理操作。 --- #### 基于 OpenVINO 的推理实现 OpenVINO 提供了一种高效的跨平台推理解决方案。下面是一段使用 OpenVINO 进行推理的 Python 示例代码[^2]。 ```python from openvino.runtime import Core # 初始化核心对象 ie_core = Core() # 加载模型文件 (.xml/.bin 或 .onnx 格式) model_path = 'path_to_model.xml' # 替换为实际路径 compiled_model = ie_core.compile_model(model=model_path, device_name="CPU") # 获取输入和输出张量的信息 input_layer = compiled_model.input(0) output_layer = compiled_model.output(0) # 准备输入数据（替换为实际图像处理后的 NumPy 数组） dummy_input = np.random.rand(*list(input_layer.shape)).astype(np.float32) # 执行推理 result = compiled_model([dummy_input])[output_layer] print(result) ``` 上述代码片段说明了如何利用 OpenVINO 配置环境、加载模型以及运行推理过程。 --- ### 总结 以上分别提供了基于 **TensorRT** 和 **OpenVINO** 的 Python 推理引擎实现方法。两者各有优劣：TensorRT 更适合 NVIDIA 平台上的高效推理；而 OpenVINO 则提供更广泛的硬件兼容性和灵活性。