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

时间:2019-05-30 03:03:58

标签: python multithreading pycuda tensorrt

我是TensorRT和CUDA的新手,我正在尝试使用TensorRT Python API实现推理服务器。我遵循end_to_end_tensorflow_mnistuff_ssd的示例,一切正常。但是,当我尝试使用引擎在多个线程中进行推理时,遇到了一些问题。所以我想知道在多线程中运行TensorRT的正确方法是什么。

这就是我尝试过的。首先,我在主线程中创建推理引擎。在辅助线程中,我使用在主线程中创建的引擎分配内存空间,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)

上面的代码产生以下错误:

pycuda._driver.LogicError: explicit_context_dependent failed: invalid device context - no currently active context?

这听起来像工作线程中没有活动的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. 这是否意味着我必须在工作线程中重建引擎,从而显着增加推理过热?
  2. 我可以在主线程和工作线程之间共享GPU上下文,这样我就不必为每个新请求创建一个新的GPU上下文了吗?如果是这样,如何在pycuda中做到这一点?

任何建议将不胜感激。谢谢!

1 个答案:

答案 0 :(得分:1)

请参阅this

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

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

并在执行GPU任务后清理:

ctx.pop()
del ctx