如何在多个线程中运行TensorRT？

Question

我是TensorRT和CUDA的新手，我正在尝试使用TensorRT PythonAPI实现推理服务器。我遵循end_to_end_tensorflow_mnist和uff_ssd示例，一切正常。但是，当我试图使用引擎在多个线程中进行推理时，我遇到了一些问题。想知道TensorRT是如何正确运行多线程的。在

这是我尝试过的。首先，我在主线程中创建推理机。在worker线程中，我使用主线程中创建的引擎分配内存空间、CUDA流和执行上下文，并进行推断：

import pycuda.autoinit  # Create CUDA context
import pycuda.driver as cuda 

# Main thread
with open(“sample.engine”, “rb”) as f, trt.Runtime(TRT_LOGGER) as runtime:
    engine = runtime.deserialize_cuda_engine(f.read())

...

# Worker thread
with engine.create_execution_context() as context:
    inputs, outputs, bindings, stream = common.allocate_buffers(engine)
    common.do_inference(context, inputs, outputs, bindings, stream)

上述代码产生以下错误：

^{pr2}$

这听起来好像在工作线程中没有活动的CUDA上下文。所以，我尝试在工作线程中手动创建CUDA上下文：

# Worker thread
from pycuda.tools import make_default_context()
cuda.init()  # Initialize CUDA
ctx = make_default_context()  # Create CUDA context

with engine.create_execution_context() as context:
    inputs, outputs, bindings, stream = common.allocate_buffers(engine)
    common.do_inference(context, inputs, outputs, bindings, stream)

ctx.pop()  # Clean up

这一次，它给了我另一个错误：

[TensorRT] ERROR: cuda/cudaConvolutionLayer.cpp (163) - Cudnn Error in execute: 7
[TensorRT] ERROR: cuda/cudaConvolutionLayer.cpp (163) - Cudnn Error in execute: 7

我知道构建器或运行时将使用与创建线程关联的GPU上下文创建。我想这个错误是因为引擎与主线程关联，但我在工作线程中使用它，所以我的问题是：

1. 这是否意味着我必须在worker线程中重建引擎，这会显著增加推理过热？在
2. 我可以在主线程和工作线程之间共享GPU上下文吗，这样我就不必为每个新请求创建一个新的GPU上下文？如果是的话，在皮库达怎么做？在

如有任何建议，我们将不胜感激。谢谢！在

Answer 1

参考this。在

如果要在多线程中进行推理，则需要修改common.py，如下所示。在触发GPU任务之前创建上下文：

dev = cuda.Device(0) // 0 is your GPU number
ctx = dev.make_context()

完成GPU任务后清理：

^{pr2}$