解决ONNX与TensorRT并行运行时CUDA资源冲突的指南

[TRT] [E] 1: [convolutionRunner.cpp::execute::391] Error Code 1: Cask (Cask convolution execution)
[TRT] [E] 1: [checkMacros.cpp::catchCudaError::272] Error Code 1: Cuda Runtime (invalid resource handle)

import cv2
import numpy as np
import pycuda.driver as cuda
import tensorrt as trt
import onnx
import onnxruntime

# 解决Numpy版本兼容性问题
np.bool = np.bool_

# 假设profliing模块可用，否则需要移除或替换
from profiling import GlobalProfTime, ProfTimer, mode_to_str

# 1. 手动初始化CUDA并创建上下文
# 移除 import pycuda.autoinit
cuda.init()
device = cuda.Device(0) # 选择第一个GPU设备
ctx = device.make_context() # 创建CUDA上下文

try:
    with GlobalProfTime('profile_tensorrt_10_000images') as t:
        with ProfTimer('TensorRT basic image profiler') as t:

            # TensorRT code
            # 加载TensorRT引擎
            TRT_ENGINE_PATH = '/app/models/buffalo_l/det_10g640x640.engine'

            # 创建运行时，并在创建引擎前激活PyCUDA上下文
            # 注意：TensorRT内部会创建自己的CUDA context，但PyCUDA的上下文需要先存在
            runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))

            # 反序列化引擎
            with open(TRT_ENGINE_PATH, 'rb') as f:
                engine_data = f.read()
                engine = runtime.deserialize_cuda_engine(engine_data)

            assert engine is not None, "TensorRT engine deserialization failed."

            # 创建执行上下文
            context = engine.create_execution_context()

            # 为输入输出分配内存
            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))
                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({'host': host_mem, 'device': device_mem, 'name': binding, 'shape': engine.get_binding_shape(binding), 'type': engine.get_binding_dtype(binding)})
                else:
                    outputs.append({'host': host_mem, 'device': device_mem, 'name': binding, 'shape': engine.get_binding_shape(binding), 'type': engine.get_binding_dtype(binding)})

            # 加载并预处理输入图片
            image_path = "/app/models/buffalo_l/image.png"
            image = cv2.imread(image_path)
            assert image is not None, f"Failed to load image from {image_path}"

            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            image = cv2.resize(image, (640, 640))
            image = image.astype(np.float32) / 255.0
            input_data = np.expand_dims(image.transpose(2, 0, 1), axis=0)

            # ONNX code
            onnx_model_path = "/app/models/buffalo_l/det_10g.onnx"
            # onnx.load(onnx_model_path) # ONNX模型加载不是必须的，InferenceSession会自动加载
            # 创建ONNX Runtime Session，指定CUDAExecutionProvider
            ort_session = onnxruntime.InferenceSession(onnx_model_path, providers=['CUDAExecutionProvider'])

            # TensorRT 推理循环
            for _ in range(1):
                with ProfTimer('TensorRT per call') as t:
                    # 复制输入数据到GPU
                    cuda.memcpy_htod_async(inputs[0]['device'], input_data.ravel(), stream)
                    # 执行推理
                    if context.execute_async(batch_size=1, bindings=bindings, stream_handle=stream.handle) == 0:
                        print("Error: Unable to launch TensorRT inference.")
                    # 将结果从GPU传回主机
                    if cuda.memcpy_dtoh_async(outputs[0]['host'], outputs[0]['device'], stream) == 0:
                        print("Error: Unable to copy results from GPU to host.")

                    result = outputs[0]['host']
                    # 同步流
                    stream.synchronize()

                    print("Inference TensorRT Results:")
                    print(result[:20])
            stream.synchronize() # 确保所有TensorRT操作完成

            # ONNX 推理循环
            for _ in range(1):
                with ProfTimer('ONNX(CUDA) per call') as t:
                    # 重新加载和预处理图像（如果需要，或使用TensorRT已加载的）
                    image_path = "/app/models/buffalo_l/image.png"
                    image = cv2.imread(image_path)
                    assert image is not None, f"Failed to load image from {image_path}"

                    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                    image = cv2.resize(image, (640, 640))
                    image = image.astype(np.float32) / 255.0
                    input_data_onnx = np.expand_dims(image.transpose(2, 0, 1), axis=0) # 使用独立的变量名避免混淆

                    # 运行ONNX推理
                    input_name = ort_session.get_inputs()[0].name
                    outputs_onnx = ort_session.run(None, {input_name: input_data_onnx})

                    print("Inference ONNX Results:")
                    print(f"{np.transpose(outputs_onnx[0][:20])}")

finally:
    # 2. 在程序结束或不再需要CUDA时，显式释放上下文
    # 这一步非常重要，确保资源被正确清理，避免潜在的内存泄漏或冲突
    if ctx:
        ctx.pop()
        ctx.detach() # 从当前线程分离上下文