TAO系列13-执行部署的TensorRT加速引擎

def _preprocess_trt(img, shape=(300, 300)):

    """TRT SSD推理前的数据前处理"""

    img = cv2.resize(img, shape)

    img = img.transpose((2, 0, 1)).astype(np.float32)

    return img

def _postprocess_trt(img, output, conf_th, output_layout):

    """TRT SSD推理后的结果的数据处理步骤."""

    img_h, img_w, _ = img.shape

    boxes, confs, clss = [], [], []

    for prefix in range(0, len(output), output_layout):

        index = int(output[prefix+0])

        conf = float(output[prefix+2])

        if conf < conf_th:

            continue

        x1 = int(output[prefix+3] * img_w)

        y1 = int(output[prefix+4] * img_h)

        x2 = int(output[prefix+5] * img_w)

        y2 = int(output[prefix+6] * img_h)

        cls = int(output[prefix+1])

        boxes.append((x1, y1, x2, y2))

        confs.append(conf)

        clss.append(cls)

    return boxes, confs, clss   # 返回标框坐标、置信度、类别

class TrtSSD(object):

# 加载自定义组建，如果TRT版本小于7.0需要额外生成flattenconcat自定义组件库

    def _load_plugins(self):

        if trt.__version__[0] < '7':

            ctypes.CDLL("ssd/libflattenconcat.so")

        trt.init_libnvinfer_plugins(self.trt_logger, '')

#加载通过Transfer Learning Toolkit生成的推理引擎

    def _load_engine(self):

        TRTbin = 'ssd/TRT_%s.bin' % self.model  #请根据实际状况自行修改

        with open(TRTbin, 'rb') as f, trt.Runtime(self.trt_logger) as runtime:

            return runtime.deserialize_cuda_engine(f.read())

#通过加载的引擎，生成可执行的上下文

    def _create_context(self):

        for binding in self.engine:

            size = trt.volume(self.engine.get_binding_shape(binding)) * \

                   self.engine.max_batch_size

##注意：这里的host_mem需要时用pagelocked memory，以免内存被释放

            host_mem = cuda.pagelocked_empty(size, np.float32)

            cuda_mem = cuda.mem_alloc(host_mem.nbytes)

            self.bindings.append(int(cuda_mem))

            if self.engine.binding_is_input(binding):

                self.host_inputs.append(host_mem)

                self.cuda_inputs.append(cuda_mem)

            else:

                self.host_outputs.append(host_mem)

                self.cuda_outputs.append(cuda_mem)

        return self.engine.create_execution_context()

# 初始化引擎

    def __init__(self, model, input_shape, output_layout=7):

        self.model = model

        self.input_shape = input_shape

        self.output_layout = output_layout

        self.trt_logger = trt.Logger(trt.Logger.INFO)

        self._load_plugins()

        self.engine = self._load_engine()

        self.host_inputs = []

        self.cuda_inputs = []

        self.host_outputs = []

        self.cuda_outputs = []

        self.bindings = []

        self.stream = cuda.Stream()

        self.context = self._create_context()

# 释放引擎，释放GPU显存，释放CUDA流

    def __del__(self):

        del self.stream

        del self.cuda_outputs

        del self.cuda_inputs

# 利用生成的可执行上下文执行推理

    def detect(self, img, conf_th=0.3):

        img_resized = _preprocess_trt(img, self.input_shape)

        np.copyto(self.host_inputs[0], img_resized.ravel())

        # 将处理好的图片从CPU内存中复制到GPU显存

        cuda.memcpy_htod_async(

            self.cuda_inputs[0], self.host_inputs[0], self.stream)

        # 开始执行推理任务

        self.context.execute_async(

            batch_size=1,

            bindings=self.bindings,

            stream_handle=self.stream.handle)

        # 将推理结果输出从GPU显存复制到CPU内存

        cuda.memcpy_dtoh_async(

            self.host_outputs[1], self.cuda_outputs[1], self.stream)

        cuda.memcpy_dtoh_async(

            self.host_outputs[0], self.cuda_outputs[0], self.stream)

        self.stream.synchronize()

        output = self.host_outputs[0]

        return _postprocess_trt(img, output, conf_th, self.output_layout)