Python中结合ONNX(CUDA Provider)与TensorRT运行报错求助

问题现象 #

同时运行基于CUDA Execution Provider的ONNX模型和TensorRT模型时，触发以下错误：

2023-11-26 11:46:35.483254243 [W:onnxruntime:, session_state.cc:1162 VerifyEachNodeIsAssignedToAnEp] Some nodes were not assigned to the preferred execution providers which may or may not have an negative impact on performance. e.g. ORT explicitly assigns shape related ops to CPU to improve perf.
2023-11-26 11:46:35.483279701 [W:onnxruntime:, session_state.cc:1164 VerifyEachNodeIsAssignedToAnEp] Rerunning with verbose output on a non-minimal build will show node assignments.
[11/26/2023-11:46:36] [TRT] [E] 1: [convolutionRunner.cpp::execute::391] Error Code 1: Cask (Cask convolution execution)
[11/26/2023-11:46:36] [TRT] [E] 1: [checkMacros.cpp::catchCudaError::272] Error Code 1: Cuda Runtime (invalid resource handle)

• 单独运行任意一个模型无错误
• 将ONNX Runtime的执行器改为CPUExecutionProvider时，也无错误

根因分析 #

错误核心是CUDA上下文冲突：

• TensorRT通过pycuda创建并管理独立的CUDA上下文
• ONNX Runtime的CUDA Execution Provider会自动初始化自己的CUDA上下文
• 两个上下文未共享，导致TensorRT的CUDA资源句柄在ORT初始化后失效，触发invalid resource handle错误

解决方案 #

1. 强制共享CUDA上下文 #

让ONNX Runtime复用pycuda已经创建的CUDA上下文，需在创建ORT会话前显式绑定当前上下文。

2. 统一CUDA流（可选优化） #

让两个模型共享同一个CUDA流，避免异步操作的同步冲突。

修改后的完整代码 #

import cv2
import numpy as np
import pycuda.driver as cuda
import pycuda.autoinit
import tensorrt as trt
np.bool = np.bool_

import onnx
import onnxruntime

from profiling import GlobalProfTime, ProfTimer, mode_to_str


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'
        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
        context = engine.create_execution_context()

        # 分配CUDA内存与创建流
        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
        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)

        # ------------------- 关键修改：绑定CUDA上下文给ONNX Runtime -------------------
        current_ctx = cuda.Context.get_current()
        ort_config = onnxruntime.SessionOptions()
        # 让ORT复用已有的CUDA上下文
        ort_config.add_session_config_entry("cuda.context_handle", str(int(current_ctx.handle)))
        # -------------------------------------------------------------------------

        # ONNX Runtime 初始化
        onnx_model_path = "/app/models/buffalo_l/det_10g.onnx"
        onnx_model = onnx.load(onnx_model_path)
        ort_session = onnxruntime.InferenceSession(onnx_model_path, sess_options=ort_config, providers=['CUDAExecutionProvider'])

        # TensorRT 推理
        for _ in range(1):
            with ProfTimer('TensorRT per call') as t:
                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.")
                cuda.memcpy_dtoh_async(outputs[0]['host'], outputs[0]['device'], stream)
                stream.synchronize()
                result = outputs[0]['host']
                print("Inference TensorRT Results:")
                print(result[:20])
        stream.synchronize()

        # ONNX Runtime 推理
        for _ in range(1):
            with ProfTimer('ONNX(CUDA) per call') as t:
                image = cv2.imread(image_path)
                assert image is not None
                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)
                input_name = ort_session.get_inputs()[0].name
                outputs = ort_session.run(None, {input_name: input_data})
                print("Inference ONNX Results:")
                print(f"{np.transpose(outputs[0][:20])}")

验证说明 #

修改后重新运行代码，invalid resource handle错误会消失，两个模型均可正常在CUDA设备上执行推理。

内容的提问来源于stack exchange，提问作者Arthur