[人工智能]构建TensorRT引擎加速MaskRCNN

环境

Jetson AGX Xavier Linux Jetpack 4.5.1
cuda 10.2
cudnn 8.0
TensorRT 7.1.3

构建引擎

有两种方案，一种是直接在代码中使用uff模型文件构建引擎进行推理，另一种是利用uff模型文件生成引擎文件engine再调用，其实道理差不多，都是要构建引擎再推理，只是一种在线构建调用，一种离线生成再调用

方案一 在线构建引擎

可参考tensorrt/sample/python/end_to_end_tensorflow_mnist/下的样例代码

from PIL import Image
import numpy as np
import os
import tensorrt as trt
import common
# You can set the logger severity higher to suppress messages (or lower to display more messages).
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

class ModelData(object):
    MODEL_FILE = 'mrcnn_nchw.uff'
    # 网络输入大小
    INPUT_SHAPE = (3, 1024, 1024)
    DTYPE = trt.float16
    # 网络输入节点名称
    INPUT_NAME ="input_image"
    # 网络输出节点名称
    OUTPUT_NAME = "mrcnn_mask/Sigmoid"

# 从uff构建引擎
def build_engine_uff(model_file):
    with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.UffParser() as parser:
        builder.max_workspace_size = 1<<30
        parser.register_input(ModelData.INPUT_NAME, ModelData.INPUT_SHAPE)
        parser.register_output(ModelData.OUTPUT_NAME)
        parser.parse(model_file, network)
        builder.build_cuda_engine(network)
        return builder.build_cuda_engine(network)

# 输入图像处理
def load_normalized_test_case(test_image, pagelocked_buffer):
    # 将图片转换为CHW
    def normalize_image(image):
        c, h, w = ModelData.INPUT_SHAPE
        image_arr = np.asarray(image.resize((w, h), Image.ANTIALIAS)).transpose([2, 0, 1]).astype(trt.nptype(ModelData.DTYPE)).ravel()
        return (image_arr / 255.0 - 0.45) / 0.225

    np.copyto(pagelocked_buffer, normalize_image(Image.open(test_image)))
    return test_image

def main():
    data_paths = 'image.jpg'
    uff_model_file = ModelData.MODEL_FILE
    # Build a TensorRT engine.
    trt.init_libnvinfer_plugins(TRT_LOGGER, '')  # 加载所有自定义的plugin
    with build_engine_uff(uff_model_file) as engine:
        inputs, outputs, bindings, stream = common.allocate_buffers(engine)
        with engine.create_execution_context() as context:
        	case_num = load_normalized_test_case(imagepath, pagelocked_buffer=inputs[0].host)
            # For more information on performing inference, refer to the introductory samples.
            # The common.do_inference function will return a list of outputs - we only have one in this case.
            # 执行推理
            output = common.do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
            box = output[0]
            mask = output[1]
            detections = box.reshape(1, 100, 6)
            mrcnn_mask = mask.reshape(1, 100, 15, 28, 28)
            print(box)
            print(mask)
            # 下面就是box和mask的后期处理了
if __name__ == '__main__':
    main()

方案二 引擎文件调用

使用trtexec可执行文件生成引擎文件，直接在终端执行命令./trtexec会打印出详细的参数说明

./trtexec --uff=/home/siu/tensorrt/bin/maskrcnn/mrcnn_nchw.uff --saveEngine=/home/siu/tensorrt/bin/maskrcnn/mrcnn_nchw.engine --workspace=8192 --uffInput=input_image,3,1024,1024 --output=mrcnn_detection,mrcnn_mask/Sigmoid --plugins=/home/siu/TensorRT/build1/out/libnvinfer_plugin.so

转换后就可以直接使用引擎文件推理了

目前这个方案碰到了一点问题，等有时间再看看了。。。。