TensorRT python多进程推理踩坑

1、TensorRT多线程推理

此时需要在每个线程内创建上下文，且在初始化ctx之后进行ctx.pop(),推理操作前，ctx.push(),运行结束前运行ctx.pop()

class detector():
    def __init__(device):
        self.ctx = cuda.Device(device).make_context()
        self.engine = self.get_engine
        self.context = self.get_context
        self.inputs, self.outputs, self.bindings, self.stream = self.allocate_buffers(self.engine)
        self.inference_fn = self.do_inference if trt.__version__[0] < '7' \
            else self.do_inference_v2
        #此处pop必须添加，实测不添加会报错
        self.ctx.pop()
    def do_inference_v2(self,context, bindings, inputs, outputs, stream):
		#此处必须添加push
        self.ctx.push()

        # Transfer input data to the GPU.
        [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
        # Run inference.
        context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
        # Transfer predictions back from the GPU.
        [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
        # Synchronize the stream
        stream.synchronize()
        # Return only the host outputs.
        return [out.host for out in outputs]
    
    def detect():
        """
        pre process
        """
        res = self.do_inference_v2(context, bindings, inputs, outputs, stream)
        """
        post process
        """
        #此处必须pop,若后处理不需要cuda和GPU，则可在postProcess之前pop
        self.ctx.pop()

2、采用multiprocessing多进程运行利用cuda库使用GPU推理时，需要设置为spawn运行

报错信息： RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the ‘spawn’ start method

#第一种方式，set_start_method只能调用一次，不然会报错
import multiprocessing as mp
mp.set_start_method('spawn')
#第二种方式
import multiprocessing as mp
ctx = mp.get_context('spawn')
q = ctx.Queue()
p = ctx.Process(target=fun,args = (para1,))

3、采用多进程进行推理时，传入参数需要是可以通过pickle序列化的，自定义对象有时会出问题

报错信息：AttributeError: Can’t pickle local object ‘_createenviron..encode’

python可以序列化的对象

None,True ，False;
整数，浮点数，复数;
字符串，字节流，字节数组;
包含可pickle对象的tuples，lists，sets和dictionaries；
定义在module顶层的函数：
定义在module顶层的内置函数；
定义在module顶层的类；
拥有__dict__()或__setstate__()的自定义类型；

解决方案

#方法一，将自定义对象定义为全局变量
#方法二，将对象中需要的结构分离出来传入进程
#方法三，使用dill包,传入参数时使用dill.dumps(obj),函数内部使用dill.loads(obj)恢复对象  ------方法三已经过检验
#方法四，在自定类中定义__getstate__和__setstate__，如果定义了这两个方法，pickle.dump() 就会调用 __getstate__() 获取序列化的对象。 类似的，__setstate__() 在反序列化时被调用。

误区一:

导入改为from multiprocessing.dummy import Queue,Process

虽然能解决报错，但这样多进程就变成多线程了

误区二：

将start()改为run()运行

能解决报错，但在多进程下run方法启动相当于直接调用函数，并没有真正意义上使用多进程，这一点我们可以通过pid看的出来。而start启动却是真正意义上调用了多进程，同样我们可以通过pid看的出来