### 基于 Python 的多 TensorRT 模型并发推理方法及实现 为了实现在 Python 中同时推理多个 TensorRT 模型，可以采用多线程或多进程的方式来管理不同的模型实例。以下是具体的解决方案： #### 1. 多线程方式 在多线程环境中，可以通过创建独立的线程来加载和运行每个 TensorRT 模型。这种方式适合轻量级的任务调度。 - **优点**: 轻便灵活，易于实现。 - **缺点**: 如果模型较大或者计算密集度高，则可能受到 GIL (Global Interpreter Lock) 的影响。 ```python import threading import tensorrt as trt from cuda import cudart class TRTInferenceThread(threading.Thread): def __init__(self, model_path, input_data): super(TRTInferenceThread, self).__init__() self.model_path = model_path self.input_data = input_data self.output_data = None def run(self): with open(self.model_path, 'rb') as f, \ trt.Runtime(trt.Logger(trt.Logger.WARNING)) as runtime: engine = runtime.deserialize_cuda_engine(f.read()) context = engine.create_execution_context() d_input = cudart.cudaMalloc(engine.get_binding_shape(0).prod() * 4)[1] d_output = cudart.cudaMalloc(engine.get_binding_shape(1).prod() * 4)[1] stream = cudart.cudaStreamCreate()[1] cudart.cudaMemcpyAsync(d_input, self.input_data.ctypes.data, engine.get_binding_shape(0).prod() * 4, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, stream) context.execute_v2([int(d_input), int(d_output)]) cudart.cudaMemcpyAsync(self.output_data.ctypes.data, d_output, engine.get_binding_shape(1).prod() * 4, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, stream) cudart.cudaStreamSynchronize(stream) cudart.cudaFree(d_input) cudart.cudaFree(d_output) cudart.cudaStreamDestroy(stream) # 创建两个线程分别执行不同模型的推理 thread1 = TRTInferenceThread('model1.trt', input_data_1) thread2 = TRTInferenceThread('model2.trt', input_data_2) thread1.start() thread2.start() thread1.join() thread2.join() ``` 上述代码展示了如何利用 `threading` 库为每个 TensorRT 模型分配单独的线程[^1]。 --- #### 2. 多进程方式 对于更复杂的场景，推荐使用多进程方案。由于每个进程拥有独立的内存空间，因此不会受 GIL 影响，更适合大规模并行任务。 - **优点**: 更高的性能和稳定性。 - **缺点**: 需要额外的资源管理和通信开销。 ```python from multiprocessing import Process, Queue import tensorrt as trt from cuda import cudart def inference_worker(model_path, input_queue, output_queue): with open(model_path, 'rb') as f, \ trt.Runtime(trt.Logger(trt.Logger.WARNING)) as runtime: engine = runtime.deserialize_cuda_engine(f.read()) context = engine.create_execution_context() while True: try: input_data = input_queue.get(timeout=1) except Exception: break d_input = cudart.cudaMalloc(engine.get_binding_shape(0).prod() * 4)[1] d_output = cudart.cudaMalloc(engine.get_binding_shape(1).prod() * 4)[1] stream = cudart.cudaStreamCreate()[1] cudart.cudaMemcpyAsync(d_input, input_data.ctypes.data, engine.get_binding_shape(0).prod() * 4, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, stream) context.execute_v2([int(d_input), int(d_output)]) output_buffer = bytearray(engine.get_binding_shape(1).prod() * 4) cudart.cudaMemcpyAsync(output_buffer, d_output, engine.get_binding_shape(1).prod() * 4, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, stream) cudart.cudaStreamSynchronize(stream) cudart.cudaFree(d_input) cudart.cudaFree(d_output) cudart.cudaStreamDestroy(stream) output_queue.put(output_buffer) input_queues = [Queue(), Queue()] output_queues = [Queue(), Queue()] processes = [ Process(target=inference_worker, args=('model1.trt', input_queues[0], output_queues[0])), Process(target=inference_worker, args=('model2.trt', input_queues[1], output_queues[1])) ] for p in processes: p.start() # 向队列发送输入数据 input_queues[0].put(input_data_1) input_queues[1].put(input_data_2) for p in processes: p.terminate() ``` 此代码片段说明了如何通过 `multiprocessing.Process` 并发运行多个 TensorRT 推理引擎[^2]。 --- #### 3. 使用 GPU 流控制优化 无论选择哪种方式，都可以进一步引入 CUDA Stream 提升效率。CUDA Streams 可以让设备上的操作异步化，从而减少等待时间。 ```python stream = cudart.cudaStreamCreate()[1] context.execute_async_v2(bindings=[int(d_input), int(d_output)], stream_handle=stream) cudart.cudaStreamSynchronize(stream) ``` 以上代码展示了一个简单的流同步机制[^3]。 --- #### 性能考量 当涉及 INT8 或 FP16 精度时，需注意量化过程是否会带来精度损失以及硬件兼容性问题。如果模型支持 PTQ/QAT，则可通过 TensorRT 自动生成高效的 Engine 文件。 ---