解决ONNX Runtime与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 pycuda.autoinit # 移除此行，手动管理CUDA上下文
import tensorrt as trt
np.bool = np.bool_ # 兼容旧版Numpy

import onnx
import onnxruntime

# 假设 profiling 模块已定义，如果不需要可删除或替换
# from profiling import GlobalProfTime, ProfTimer, mode_to_str

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

# 2. 将创建的上下文推入当前线程的CUDA上下文堆栈
ctx.push()

# 使用try...finally块确保上下文在程序结束时被正确弹出和清理
try:
    # with GlobalProfTime('profile_tensorrt_10_000images') as t: # 性能分析，如果不需要可注释
    # with ProfTimer('TensorRT basic image profiler') as t: # 性能分析，如果不需要可注释

    # --- TensorRT 模型初始化部分 ---
    TRT_ENGINE_PATH = '/app/models/buffalo_l/det_10g640x640.engine'

    # 创建TensorRT运行时日志器
    runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))

    # 反序列化TensorRT引擎
    with open(TRT_ENGINE_PATH, 'rb') as f:
        engine_data = f.read()
        engine = runtime.deserialize_cuda_engine(engine_data)
    assert engine is not None, "Failed to deserialize TensorRT engine."

    # 创建TensorRT执行上下文
    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)})

    # 加载并预处理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_trt = np.expand_dims(image.transpose(2, 0, 1), axis=0)

    # --- ONNX Runtime 模型初始化部分 ---
    onnx_model_path = "/app/models/buffalo_l/det_10g.onnx"
    # 加载ONNX模型（可选，onnxruntime.InferenceSession会自动加载）
    # onnx_model = onnx.load(onnx_model_path)

    # 创建ONNX Runtime会话，并指定CUDAExecutionProvider
    # 注意：ort_session的创建必须在CUDA上下文被推入后进行
    ort_session = onnxruntime.InferenceSession(onnx_model_path, providers=['CUDAExecutionProvider'])

    # --- TensorRT 模型推理部分 ---
    print("\n--- Running TensorRT Inference ---")
    for _ in range(1):
        # with ProfTimer('TensorRT per call') as t: # 性能分析，如果不需要可注释
        # 将输入数据从主机内存拷贝到GPU设备内存
        cuda.memcpy_htod_async(inputs[0]['device'], input_data_trt.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_trt = outputs[0]['host']
        # 同步CUDA流
        stream.synchronize()

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

    # --- ONNX Runtime 模型推理部分 ---
    print("\n--- Running ONNX (CUDA) Inference ---")
    for _ in range(1):
        # with ProfTimer('ONNX(CUDA) per call') as t: # 性能分析，如果不需要可注释
        # 重新加载并预处理ONNX模型输入图像（如果需要，否则可复用input_data_trt）
        image_path_onnx = "/app/models/buffalo_l/image.png"
        image_onnx = cv2.imread(image_path_onnx)
        assert image_onnx is not None, f"Failed to load image from {image_path_onnx}"
        image_onnx = cv2.cvtColor(image_onnx, cv2.COLOR_BGR2RGB)
        image_onnx = cv2.resize(image_onnx, (640, 640))
        image_onnx = image_onnx.astype(np.float32) / 255.0
        input_data_onnx = np.expand_dims(image_onnx.transpose(2, 0, 1), axis=0)

        # 获取ONNX模型的输入名称
        input_name_onnx = ort_session.get_inputs()[0].name
        # 运行ONNX推理
        outputs_onnx = ort_session.run(None, {input_name_onnx: input_data_onnx})

        print("Inference ONNX Results (first 20 elements):")
        # 注意：ONNX输出的形状可能与TensorRT不同，这里仅打印前20个元素
        print(f"{np.transpose(outputs_onnx[0][:20])}")

finally:
    # 3. 确保上下文被弹出，无论是否发生异常
    ctx.pop()
    # 4. 分离并销毁上下文以释放资源
    ctx.detach()

    # 清理TensorRT相关资源
    if 'context' in locals() and context:
        context.destroy()
    if 'engine' in locals() and engine:
        engine.destroy()
    if 'runtime' in locals() and runtime:
        del runtime # runtime对象没有destroy方法，直接删除引用

    # ONNX Runtime session通常在对象销毁时自动释放资源，无需显式操作
    # 如果需要，也可以 del ort_session
    print("\nCUDA Context and resources cleaned up.")