Tensorrt实现solov2加速

SOLO简介

solo系列网络是由Xinlong Wang提出的单阶段实例分割网络。其搭建在mmdetection库中。solov2主干网络如下图所示：

其在COCO数据集上获得了较高的AP，并且由于其单阶段实例分割的特点，方面直接实现端到端的部署。因此，我们可以采用pytorch转onnx再转tensorrt的方式实现solo的半精度、int8加速，达到实时运行的目的。

环境依赖

我们采用TensorRT7.1的部署环境（TensorRT 7.0的instance normalization层存在bug，造成推理结果不正确，参见gihub该问题的讨论）。具体测试环境如下：

	Ubuntu 18.04
	opencv 4
	CUDA 10.1
	TensorRT 7.1.3
	pytorch 1.3

快速开始

以下代码展示了如何利用tensorrt的python接口快速进行pytorch到onnx模型的转换，大家也可以去我的github上下载。

转换pytorch模型

根据以下脚本将pytorch模型转化为onnx模型：

import argparse
import mmcv
import torch
from mmcv.runner import load_checkpoint
import torch.nn.functional as F
from mmdet.models import build_detector
import cv2
import torch.onnx as onnx
import numpy as np
import torch.nn as nn


numclass=80

def points_nms(heat, kernel=2):
    # kernel must be 2
    hmax = nn.functional.max_pool2d(
        heat, (kernel, kernel), stride=1, padding=1)
    keep = (hmax[:, :, :-1, :-1] == heat).float()
    return heat * keep


def fpn_forward(self, inputs):
    assert len(inputs) == len(self.in_channels)

    # build laterals
    laterals = [
            lateral_conv(inputs[i + self.start_level])
            for i, lateral_conv in enumerate(self.lateral_convs)
    ]
    # build top-down path
    used_backbone_levels = len(laterals)
    for i in range(used_backbone_levels - 1, 0, -1):

        sh = torch.tensor(laterals[i].shape)
        laterals[i - 1] += F.interpolate(
            laterals[i], size=(sh[2]*2,sh[3]*2), mode='nearest')

    outs = [
        self.fpn_convs[i](laterals[i]) for i in range(used_backbone_levels)
    ]
    # part 2: add extra levels
    if self.num_outs > len(outs):
        if not self.add_extra_convs:
            for i in range(self.num_outs - used_backbone_levels):
                outs.append(F.max_pool2d(outs[-1], 1, stride=2))
        # add conv layers on top of original feature maps (RetinaNet)
        else:
            if self.extra_convs_on_inputs:
                orig = inputs[self.backbone_end_level - 1]
                outs.append(self.fpn_convs[used_backbone_levels](orig))
            else:
                outs.append(self.fpn_convs[used_backbone_levels](outs[-1]))
            for i in range(used_backbone_levels + 1, self.num_outs):
                if self.relu_before_extra_convs:
                    outs.append(self.fpn_convs[i](F.relu(outs[-1])))
                else:
                    outs.append(self.fpn_convs[i](outs[-1]))
    return tuple(outs)


def forward_single(self, x, idx, eval=False, upsampled_size=None): # bbox head
    ins_kernel_feat = x
    y_range = np.linspace(-1, 1, ins_kernel_feat.shape[-1],dtype=np.float32)#h
    x_range = np.linspace(-1, 1, ins_kernel_feat.shape[-2],dtype=np.float32)#w
    x, y = np.meshgrid(y_range, x_range)
    y = y[None][None]
    x = x[None][None]
    coord_feat =np.concatenate([x, y], 1)

    coord_feat__ = torch.tensor(coord_feat)

    seg_num_grid = self.seg_num_grids[idx]
    cate_feat = F.interpolate(ins_kernel_feat, size=seg_num_grid, mode='bilinear')
    #ins_kernel_feat = torch.cat([ins_kernel_feat, coord_feat], 1)
    kernel_feat = torch.cat([ins_kernel_feat, coord_feat__], 1)
    # kernel branch
    kernel_feat = F.interpolate(kernel_feat, size=seg_num_grid, mode='bilinear')

    kernel_feat = kernel_feat.contiguous()
    for i, kernel_layer in enumerate(self.kernel_convs):

        kernel_feat = kernel_layer.conv(kernel_feat)

        num_group = torch.tensor(kernel_layer.gn.num_groups)
        sh = torch.tensor(kernel_feat.shape)

        kernel_feat = kernel_feat.view(1,num_group,-1)
        insta_weight = torch.ones(num_group)
        insta_bias = torch.zeros(num_group)
        kernel_feat = F.instance_norm(kernel_feat,weight=insta_weight)
        kernel_feat = kernel_feat.view(sh[0],sh[1],sh[2],sh[3])

        gn_weight = kernel_layer.gn.weight.data.view(1,-1,1,1)

        gn_bias = kernel_layer.gn.bias.data.view(1,-1,1,1)
        kernel_feat = gn_weight*kernel_feat
        kernel_feat = gn_bias+kernel_feat

        kernel_feat = F.relu(kernel_feat)
        #kernel_feat = kernel_layer(kernel_feat)
    kernel_pred = self.solo_kernel(kernel_feat)

        # cate branch
    cate_feat = cate_feat.contiguous()
    for i, cate_layer in enumerate(self.cate_convs):

        cate_feat = cate_layer.conv(cate_feat)

        num_group = torch.tensor(cate_layer.gn.num_groups)
        sh = torch.tensor(cate_feat.shape)
        cate_feat = cate_feat.view(1,num_group,-1)
        
        insta_weight = torch.ones(num_group)
        insta_bias = torch.zeros(num_group)
        cate_feat = F.instance_norm(cate_feat,weight=insta_weight)
        cate_feat = cate_feat.view(sh[0],sh[1],sh[2],sh[3])

        gn_weight = cate_layer.gn.weight.data.view(1,-1,1,1)

        gn_bias = cate_layer.gn.bias.data.view(1,-1,1,1)
        cate_feat = gn_weight*cate_feat
        cate_feat = gn_bias+cate_feat

        cate_feat = F.relu(cate_feat)
        #cate_feat = cate_layer(cate_feat)

    cate_pred = self.solo_cate(cate_feat)

    if eval:
       cate_pred = points_nms(cate_pred.sigmoid(), kernel=2).permute(0, 2, 3, 1)
    return cate_pred, kernel_pred


def split_feats(self, feats):
    sh1 = torch.tensor(feats[0].shape)
    sh2 = torch.tensor(feats[3].shape)

    return (F.interpolate(feats[0], size=(int(sh1[2]*0.5),int(sh1[3]*0.5)), mode='bilinear'), #从da到xiao
            feats[1],
            feats[2],
            feats[3],
            F.interpolate(feats[4], size=(sh2[2],sh2[3]), mode='bilinear'))


def reshap_gn_mask_nead(layer,inputs):

    inputs = layer.conv(inputs)

    num_group = torch.tensor(layer.gn.num_groups)
    sh = torch.tensor(inputs.shape)

    inputs = inputs.view(1,num_group,-1)
        
    insta_weight = torch.ones(num_group)
    insta_bias = torch.zeros(num_group)
    inputs = F.instance_norm(inputs,weight=insta_weight)
    inputs = inputs.view(sh[0],sh[1],sh[2],sh[3])

    gn_weight = layer.gn.weight.data.view(1,-1,1,1)

    gn_bias = layer.gn.bias.data.view(1,-1,1,1)
    inputs = gn_weight*inputs
    inputs = gn_bias+inputs

    outputs = F.relu(inputs)

    return outputs


def forward(self, inputs): #mask head
    feature_add_all_level = reshap_gn_mask_nead(self.convs_all_levels[0].conv0,inputs[0])
 
    x = reshap_gn_mask_nead(self.convs_all_levels[1].conv0,inputs[1])
    sh = torch.tensor(x.shape)
    feature_add_all_level += F.interpolate(x, size=(sh[2]*2,sh[3]*2), mode='bilinear')

    x = reshap_gn_mask_nead(self.convs_all_levels[2].conv0,inputs[2])
    sh = torch.tensor(x.shape)
    x = F.interpolate(x, size=(sh[2]*2,sh[3]*2), mode='bilinear')
    x = reshap_gn_mask_nead(self.convs_all_levels[2].conv1,x)
    sh = torch.tensor(x.shape)
    feature_add_all_level += F.interpolate(x, size=(sh[2]*2,sh[3]*2), mode='bilinear')

    y_range = np.linspace(-1, 1, inputs[3].shape[-1],dtype=np.float32)#h
    x_range = np.linspace(-1, 1, inputs[3].shape[-2],dtype=np.float32)#w
    x, y = np.meshgrid(y_range, x_range)
    y = y[None][None]
    x = x[None][None]
    coord_feat =np.concatenate([x, y], 1)
    coord_feat__ = torch.tensor(coord_feat)
    input_p = torch.cat([inputs[3], coord_feat__], 1)
    x = reshap_gn_mask_nead(self.convs_all_levels[3].conv0,input_p)
    sh = torch.tensor(x.shape)
    x = F.interpolate(x, size=(sh[2]*2,sh[3]*2), mode='bilinear')
    x = reshap_gn_mask_nead(self.convs_all_levels[3].conv1,x)
    sh = torch.tensor(x.shape)
    x = F.interpolate(x, size=(sh[2]*2,sh[3]*2), mode='bilinear')
    x = reshap_gn_mask_nead(self.convs_all_levels[3].conv2,x)
    sh = torch.tensor(x.shape)
    feature_add_all_level += F.interpolate(x, size=(sh[2]*2,sh[3]*2), mode='bilinear')

    feature_pred = reshap_gn_mask_nead(self.conv_pred[0],feature_add_all_level)

    return feature_pred


def main_forward(self,x):
    x = self.extract_feat(x)
    outs = self.bbox_head(x, eval=True)
    mask_feat_pred = self.mask_feat_head(
               x[self.mask_feat_head.
                 start_level:self.mask_feat_head.end_level + 1])
    cate_pred_list = [outs[0][i].view(-1, numclass) for i in range(5)]
    kernel_pred_list = [
    outs[1][i].squeeze(0).permute(1, 2, 0).view(-1, 256) for i in range(5)]

    cate_pred_list = torch.cat(cate_pred_list, dim=0)
    kernel_pred_list = torch.cat(kernel_pred_list, dim=0)

    return (cate_pred_list,kernel_pred_list,mask_feat_pred)


def parse_args():
    parser = argparse.ArgumentParser(description='get solo onnx model')
    parser.add_argument('--config', help='test config file path')
    parser.add_argument('--checkpoint', help='checkpoint file')
    parser.add_argument('--outputname',help="output name")
    parser.add_argument('--numclass', type=int,default=80)
    parser.add_argument('--inputh', type=int,default=800)
    parser.add_argument('--inputw', type=int,default=1344)
    args = parser.parse_args()

    return args


def main():
    args = parse_args()
    cfg = mmcv.Config.fromfile(args.config)
    cfg.model.pretrained = None
    cfg.data.test.test_mode = True
    global numclass
    # build the model and load checkpoint
    model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg)

    from types import MethodType
    model.bbox_head.forward_single = MethodType(forward_single,model.bbox_head)
    model.bbox_head.split_feats = MethodType(split_feats,model.bbox_head)
    model.mask_feat_head.forward = MethodType(forward,model.mask_feat_head)
    model.neck.forward = MethodType(fpn_forward, model.neck)

    img = torch.randn(1,3,args.inputh,args.inputw)
    model.forward = MethodType(main_forward,model)

    outs = model(img)
    print(len(outs))
    checkpoint = load_checkpoint(model, args.checkpoint, map_location='cpu')
    outputname = ["output1","output2","output3"]
    onnx.export(model,img,args.outputname,verbose=True,opset_version=10,input_names=["input"],output_names=outputname)


if __name__ == '__main__':
    main()

生成FP16模型并执行推理

第一版程序，使用pycuda

import pycuda.driver as cuda
import pycuda.autoinit
import cv2
import os
import numpy as np
import tensorrt as trt
import time
import argparse
import torch
import torch.nn.functional as F
import numpy as np


seg_num_grids = [40,36,24,16,12]
self_strides = [8,8,16,32,32]
score_thr = 0.1
mask_thr = 0.5
max_per_img = 100
class_names = [] # 输入你模型要预测类的名字


class HostDeviceMem(object):
    def __init__(self, host_mem, device_mem):
        self.host = host_mem
        self.device = device_mem

    def __str__(self):
        return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)

    def __repr__(self):
        return self.__str__()


class Preprocessimage(object):
    def __init__(self,inszie):
        self.inszie = inszie

    def process(self,image_path):
        start = time.time()
        image = cv2.imread(image_path) # bgr rgb
        image = cv2.cvtColor(image,cv2.COLOR_BGR2RGB)
        H,W,_ = image.shape

        img_metas = dict()
        image = cv2.resize(image,self.inszie) # resize
        img_metas["img_shape"] = image.shape
        image_raw = cv2.cvtColor(image,cv2.COLOR_RGB2BGR)

        image = image.transpose([2,0,1]) # chw
        image = np.expand_dims(image,axis=0) # nchw
        image = np.array(image,dtype=np.float32,order="C")
        print("preprocess time {:.3f} ms".format((time.time()-start)*1000))
        return image,image_raw,img_metas


def get_engine(onnx_path,engine_path,TRT_LOGGER,mode="fp16"):
    # 如果有engine直接用，否则构建新的engine
    def build_engine():
        EXPLICIT_BATCH = 1<<(int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
        with trt.Builder(TRT_LOGGER) as builder,\
            builder.create_network(EXPLICIT_BATCH) as network,\
            trt.OnnxParser(network,TRT_LOGGER) as parser:
            builder.max_workspace_size = 1<<30
            builder.max_batch_size = 1
            if mode=="fp16":
                builder.fp16_mode = True
            if not os.path.exists(onnx_path):
                print("onnx file {} not found".format(onnx_path))
                exit(0)
            print("loading onnx file {} .....".format(onnx_path))
            with open(onnx_path,'rb') as model:
                print("Begining parsing....")
                parser.parse(model.read())
            print("completed parsing")
            print("Building an engine from file {}".format(onnx_path))
            network.get_input(0).shape = [1,3,800,1344]
            engine = builder.build_cuda_engine(network)

            print("completed build engine")
            with open(engine_path,"wb") as f:
                f.write(engine.serialize())
            return engine
    if os.path.exists(engine_path):
        print("loading engine file {} ...".format(engine_path))
        with open(engine_path,"rb") as f,\
            trt.Runtime(TRT_LOGGER) as runtime:
            return runtime.deserialize_cuda_engine(f.read())
    else:
        return build_engine()


def allocate_buffers(engine):
    inputs = []
    outputs = []
    bindings = []
    stream = cuda.Stream()

    for binding in engine:
        size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
        dtype = trt.nptype(engine.get_binding_dtype(binding))
        host_mem = cuda.pagelocked_empty(size,dtype)
        device_mem = cuda.mem_alloc(host_mem.nbytes)

        bindings.append(int(device_mem))

        if engine.binding_is_input(binding):
            inputs.append(HostDeviceMem(host_mem,device_mem))
        else:
            outputs.append(HostDeviceMem(host_mem,device_mem))

    return inputs,outputs,bindings,stream


def do_inference(context,bindings,inputs,outputs,stream,batch_size=1):
    [cuda.memcpy_htod_async(inp.device,inp.host,stream) for inp in inputs]

    context.execute_async_v2(bindings=bindings,stream_handle=stream.handle)

    [cuda.memcpy_dtoh_async(out.host,out.device,stream) for out in outputs]

    stream.synchronize()

    return [out.host for out in outputs]


def matrix_nms(seg_masks, cate_labels, cate_scores, kernel='gaussian', sigma=2.0, sum_masks=None):
    """Matrix NMS for multi-class masks.

    Args:
        seg_masks (Tensor): shape (n, h, w)
        cate_labels (Tensor): shape (n), mask labels in descending order
        cate_scores (Tensor): shape (n), mask scores in descending order
        kernel (str):  'linear' or 'gauss'
        sigma (float): std in gaussian method
        sum_masks (Tensor): The sum of seg_masks

    Returns:
        Tensor: cate_scores_update, tensors of shape (n)
    """
    n_samples = len(cate_labels)
    if n_samples == 0:
        return []
    if sum_masks is None:
        sum_masks = seg_masks.sum((1, 2)).float()
    seg_masks = seg_masks.reshape(n_samples, -1).float()
    # inter.
    inter_matrix = torch.mm(seg_masks, seg_masks.transpose(1, 0))
    # union.
    sum_masks_x = sum_masks.expand(n_samples, n_samples)
    # iou.
    iou_matrix = (inter_matrix / (sum_masks_x + sum_masks_x.transpose(1, 0) - inter_matrix)).triu(diagonal=1)
    # label_specific matrix.
    cate_labels_x = cate_labels.expand(n_samples, n_samples)
    label_matrix = (cate_labels_x == cate_labels_x.transpose(1, 0)).float().triu(diagonal=1)

    # IoU compensation
    compensate_iou, _ = (iou_matrix * label_matrix).max(0)
    compensate_iou = compensate_iou.expand(n_samples, n_samples).transpose(1, 0)

    # IoU decay
    decay_iou = iou_matrix * label_matrix

    # matrix nms
    if kernel == 'gaussian':
        decay_matrix = torch.exp(-1 * sigma * (decay_iou ** 2))
        compensate_matrix = torch.exp(-1 * sigma * (compensate_iou ** 2))
        decay_coefficient, _ = (decay_matrix / compensate_matrix).min(0)
    elif kernel == 'linear':
        decay_matrix = (1-decay_iou)/(1-compensate_iou)
        decay_coefficient, _ = decay_matrix.min(0)
    else:
        raise NotImplementedError

    # update the score.
    cate_scores_update = cate_scores * decay_coefficient
    return cate_scores_update


def get_seg_single(cate_preds,
                seg_preds,
                kernel_preds,
                img_metas):

    img_shape = img_metas['img_shape']

    # overall info.
    h, w, _ = img_shape

    featmap_size = seg_preds.size()[-2:]
    upsampled_size_out = (featmap_size[0] * 4, featmap_size[1] * 4) #seg # 1344,800

    # process.
    inds = (cate_preds > score_thr)
    cate_scores = cate_preds[inds]
    if len(cate_scores) == 0:
        return None

    # cate_labels & kernel_preds
    inds = inds.nonzero()
    cate_labels = inds[:, 1]
    kernel_preds = kernel_preds[inds[:, 0]] # 选择cate大于阈值对应的kernel

    # trans vector.
    size_trans = cate_labels.new_tensor(seg_num_grids).pow(2).cumsum(0) # tensor([1600, 2896, 3472, 3728, 3872])
    strides = kernel_preds.new_ones(size_trans[-1]) #  [1,1,1,1,....,1] # 3872 所有的s*s累加

    n_stage = len(seg_num_grids) # 5
    strides[:size_trans[0]] *= self_strides[0] # [8,8,8,8......,8] 前1600乘8
    for ind_ in range(1, n_stage): #2,3,4,5
        strides[size_trans[ind_-1]:size_trans[ind_]] *= self_strides[ind_] # self.strides[8, 8, 16, 32, 32]
    strides = strides[inds[:, 0]] # 选择前坐标

    # mask encoding.
    I, N = kernel_preds.shape #( 选出的kernel,256)
    kernel_preds = kernel_preds.view(I, N, 1, 1) # (out_channels，in_channe/groups，H，W)
    seg_preds = F.conv2d(seg_preds, kernel_preds, stride=1).squeeze(0).sigmoid() #(选出的kernel,h,w)
    # mask.
    seg_masks = seg_preds > mask_thr
    sum_masks = seg_masks.sum((1, 2)).float()

    # filter.
    keep = sum_masks > strides # 大于相对应的stride
    if keep.sum() == 0:
        return None

    seg_masks = seg_masks[keep, ...]
    seg_preds = seg_preds[keep, ...]
    sum_masks = sum_masks[keep]
    cate_scores = cate_scores[keep]
    cate_labels = cate_labels[keep]

    # mask scoring.
    seg_scores = (seg_preds * seg_masks.float()).sum((1, 2)) / sum_masks
    cate_scores *= seg_scores

    # sort and keep top nms_pre
    sort_inds = torch.argsort(cate_scores, descending=True)
    if len(sort_inds) > max_per_img:
        sort_inds = sort_inds[:max_per_img]
    seg_masks = seg_masks[sort_inds, :, :]
    seg_preds = seg_preds[sort_inds, :, :]
    sum_masks = sum_masks[sort_inds]
    cate_scores = cate_scores[sort_inds]
    cate_labels = cate_labels[sort_inds]

    # Matrix NMS
    cate_scores = matrix_nms(seg_masks, cate_labels, cate_scores,
                                    kernel='gaussian',sigma=2., sum_masks=sum_masks)

    if seg_preds.shape[0]==1:
        seg_preds = cv2.resize(seg_preds.permute(1,2,0).numpy(),
                           (upsampled_size_out[1],upsampled_size_out[0]))[:,:,None].transpose(2,0,1)
    else:
        seg_preds = cv2.resize(seg_preds.permute(1,2,0).numpy(),
                           (upsampled_size_out[1],upsampled_size_out[0])).transpose(2,0,1)
    seg_masks = seg_masks > mask_thr
    return seg_masks, cate_labels, cate_scores


def vis_seg(image_raw,result,score_thresh,output):
    img_show = image_raw
    seg_show1 = img_show.copy()
    seg_show = img_show.copy()
    if result==None:
        cv2.imwrite(output,seg_show1)
    else:

        seg_label = result[0]
        seg_label = seg_label.astype(np.uint8)
        cate_label = result[1]
        cate_label = cate_label.numpy()
        score = result[2].numpy()

        vis_inds = score > score_thresh
        seg_label = seg_label[vis_inds]
        num_mask = seg_label.shape[0]
        cate_label = cate_label[vis_inds]
        cate_score = score[vis_inds]

        mask_density = []
        for idx in range(num_mask):
            cur_mask = seg_label[idx, :, :]

            mask_density.append(cur_mask.sum())
        orders = np.argsort(mask_density)
        seg_label = seg_label[orders]
        cate_label = cate_label[orders]
        cate_score = cate_score[orders]

    
        for idx in range(num_mask):
            idx = -(idx + 1)
            cur_mask = seg_label[idx, :, :]

            if cur_mask.sum() == 0:
                continue
            color_mask = (np.random.randint(0,255),np.random.randint(0,255),np.random.randint(0,255))
            contours,_ = cv2.findContours(cur_mask,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_NONE)
            cv2.drawContours(seg_show,contours,-1,color_mask,-1)
            cur_cate = cate_label[idx]
            label_text = class_names[cur_cate]
            x1,y1,w,h = cv2.boundingRect(cur_mask)
            x2 = x1+w
            y2 = y1+h
            vis_pos = (max(int(x1)-10,0),int(y1))
            cv2.rectangle(seg_show,(x1,y1),(x2,y2),(0,0,0),thickness=2)
            cv2.putText(seg_show,label_text,vis_pos,cv2.FONT_HERSHEY_COMPLEX,1,(0,0,0))
            seg_show1 = cv2.addWeighted(seg_show,0.7,img_show,0.5,0)
        cv2.imwrite(output,seg_show1)


def main():

    args = argparse.ArgumentParser(description="trt pose predict")
    args.add_argument("--onnx_path",type=str,default="dense121.onnx")
    args.add_argument("--engine_path",type=str,default="dense121fp16.trt")
    args.add_argument("--image_path",type=str)
    args.add_argument("--mode",type=str,default="fp16")
    args.add_argument("--output",type=str,default="result.png")
    args.add_argument("--classes", type=int, default=80)
    args.add_argument("--score_thr", type=float, default=0.3)
    opt = args.parse_args()

    insize = (1344,800)

    output_shape = [(1, 256, 200, 336),(3872,opt.classes),(3872,256)]
    TRT_LOGGER = trt.Logger()
    preprocesser = Preprocessimage(insize)

    image, image_raw,img_metas = preprocesser.process(opt.image_path)

    with get_engine(opt.onnx_path,opt.engine_path,TRT_LOGGER,opt.mode) as engine, \
        engine.create_execution_context() as context:
        inputs,outputs,bindings,stream = allocate_buffers(engine)

        inputs[0].host = image
        start = time.time()
        trt_outputs = do_inference(context,bindings,inputs,outputs,stream)
        end = time.time()
        print("inference time {:.3f} ms".format((end-start)*1000))
    start = time.time()
    trt_outputs = [output.reshape(shape) for output ,shape in zip(trt_outputs,output_shape)]
    trt_outputs = [torch.tensor(output) for output in trt_outputs]

    cate_pred = trt_outputs[1]
    kernel_pred = trt_outputs[2]
    seg_pred = trt_outputs[0]

    with torch.no_grad():
        result = get_seg_single(cate_pred,kernel_pred,seg_pred,img_metas)
    vis_seg(image_raw,result,opt.score_thr,opt.output)
    print("post time {:.3f} ms".format((end - start) * 1000))


if __name__=="__main__":
    main()

第二版程序，不使用pycuda

#!/usr/bin/env python3

import cv2
import os
import numpy as np
import tensorrt as trt
import torch.nn as nn
import torch.nn.functional as F
from torchvision.transforms import Normalize
import time
import argparse
import torch
import torch.nn as nn


seg_num_grids = [40, 36, 24, 16, 12]
self_strides = [8, 8, 16, 32, 32]
score_thr = 0.1
mask_thr = 0.5
update_thr = 0.05
nms_pre =500
max_per_img = 100
class_num = 1000 # ins
colors = [(np.random.random((1, 3)) * 255).tolist()[0] for i in range(class_num)]
class_names = ["person", "bicycle", "car", "motorcycle", "airplane", "bus",
               "train", "truck", "boat", "traffic_light", "fire_hydrant",
               "stop_sign", "parking_meter", "bench", "bird", "cat", "dog",
               "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe",
               "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
               "skis", "snowboard", "sports_ball", "kite", "baseball_bat",
               "baseball_glove", "skateboard", "surfboard", "tennis_racket",
               "bottle", "wine_glass", "cup", "fork", "knife", "spoon", "bowl",
               "banana", "apple", "sandwich", "orange", "broccoli", "carrot",
               "hot_dog", "pizza", "donut", "cake", "chair", "couch",
               "potted_plant", "bed", "dining_table", "toilet", "tv", "laptop",
               "mouse", "remote", "keyboard", "cell_phone", "microwave",
               "oven", "toaster", "sink", "refrigerator", "book", "clock",
               "vase", "scissors", "teddy_bear", "hair_drier", "toothbrush"]


def torch_dtype_from_trt(dtype):
    if dtype == trt.bool:
        return torch.bool
    elif dtype == trt.int8:
        return torch.int8
    elif dtype == trt.int32:
        return torch.int32
    elif dtype == trt.float16:
        return torch.float16
    elif dtype == trt.float32:
        return torch.float32
    else:
        raise TypeError('%s is not supported by torch' % dtype)


def torch_device_to_trt(device):
    if device.type == torch.device('cuda').type:
        return trt.TensorLocation.DEVICE
    elif device.type == torch.device('cpu').type:
        return trt.TensorLocation.HOST
    else:
        return TypeError('%s is not supported by tensorrt' % device)


def torch_device_from_trt(device):
    if device == trt.TensorLocation.DEVICE:
        return torch.device('cuda')
    elif device == trt.TensorLocation.HOST:
        return torch.device('cpu')
    else:
        return TypeError('%s is not supported by torch' % device)


class Preprocessimage(object):
    def __init__(self,inszie):
        self.inszie = (inszie[3],inszie[2])
		self.Normalize = Normalize(mean=[0.485,0.456,0.406],std=[0.229,0.224,0.225] ) 
		
    def process(self,image_path):
        start = time.time()
        image = cv2.imread(image_path)#[...,::-1] # bgr rgb
        image = cv2.cvtColor(image,cv2.COLOR_BGR2RGB)
        img_metas = dict()
        H,W,_ = image.shape

        image = cv2.resize(image,self.inszie) #10ms
        new_H,new_W,_ = image.shape
        img_metas["img_shape"] = image.shape
        image_raw =  cv2.cvtColor(image,cv2.COLOR_RGB2BGR)
        image = torch.form_numpy(image).float().cuda()
        image = image.permute(2,0,1) # chw
        image = self.Normalize(image/255.)
        image = image.unsqueeze(0)

        return image,image_raw,img_metas


class TRT_model(nn.Module):
    def __init__(self,
                input_size,
                onnx_path,
                engine_path,
                mode="fp16"):
        super(TRT_model, self).__init__()
        self._register_state_dict_hook(TRT_model._on_state_dict)
        self.TRT_LOGGER = trt.Logger()
        self.onnx_path = onnx_path
        self.engine_path = engine_path
        self.input_size = input_size
        self.mode = mode
        
        if os.path.exists(engine_path):
            print("loading engine file {} ...".format(engine_path))
            trt.init_libnvinfer_plugins(self.TRT_LOGGER,"")
            with open(engine_path,"rb") as f,\
                trt.Runtime(self.TRT_LOGGER) as runtime:
                    self.engine = runtime.deserialize_cuda_engine(f.read())
        else:
            self.engine = self.build_engine()

        self.context = self.engine.create_execution_context()  

    def _on_state_dict(self, state_dict, prefix, local_metadata):
        state_dict[prefix + 'engine'] = bytearray(self.engine.serialize())

    def build_engine(self):
        EXPLICIT_BATCH = 1<<(int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
        with trt.Builder(self.TRT_LOGGER) as builder,\
            builder.create_network(EXPLICIT_BATCH) as network,\
            trt.OnnxParser(network,self.TRT_LOGGER) as parser:

            builder.max_workspace_size =1<<20
            builder.max_batch_size = 1
            if self.mode =="fp16":
                print("build fp16 mode")
                builder.fp16_mode = True
            if not os.path.exists(self.onnx_path):
                print("onnx file {} not found".format(self.onnx_path))
                exit(0)
            print("loading onnx file {} .....".format(self.onnx_path))

            with open(self.onnx_path,'rb') as model:
                print("Begining parsing....")
                parser.parse(model.read())
            print("completed parsing")
            print("Building an engine from file {}".format(self.onnx_path))

            network.get_input(0).shape = self.input_size 
            engine = builder.build_cuda_engine(network)

            print("completed build engine")
            with open(self.engine_path,"wb") as f:
                f.write(engine.serialize())
            return engine

    def forward(self,inputs):
        #start = time.time()
        bindngs = [None]*(1+3)
        bindngs[0]= inputs.contiguous().data_ptr()

        outputs = [None]*3
        for i in range(1,4):
            output_shape = tuple(self.context.get_binding_shape(i))
            dtype = torch_dtype_from_trt(self.engine.get_binding_dtype(i))
            device = torch_device_from_trt(self.engine.get_location(i))

            output = torch.empty(size=output_shape,dtype=dtype,device=device)
            outputs[i-1] = output
            bindngs[i] = output.data_ptr()
        
        self.context.execute_async_v2(bindngs,
                torch.cuda.current_stream().cuda_stream)

        cate_preds = outputs[1]
        kernel_preds = outputs[2]
        seg_pred = outputs[0]

        result = get_seg_single(cate_preds,kernel_preds,seg_pred)
        #print("conv2d time {:.3f} ms".format((time.time() - start) * 1000))
        return result


def get_seg_single(cate_preds,
                    kernel_preds,
                    seg_preds
                    ):
    # process.
    inds = (cate_preds > 0.1) # 选出类
    cate_scores = cate_preds[inds]
    if len(cate_scores) == 0:
        return None

    # cate_labels & kernel_preds
    inds = inds.nonzero()
    cate_labels = inds[:, 1]
    kernel_preds = kernel_preds[inds[:, 0]]

    # trans vector.
    #print(seg_num_grids)
    size_trans = cate_labels.new_tensor(seg_num_grids).pow(2).cumsum(0)
    strides = kernel_preds.new_ones(size_trans[-1])# 3872个1

    n_stage = len(seg_num_grids)
    strides[:size_trans[0]] *= self_strides[0]
    for ind_ in range(1, n_stage):
        strides[size_trans[ind_-1]:size_trans[ind_]] *= self_strides[ind_]
    strides = strides[inds[:, 0]] # [8.8.8.]
    
    # mask encoding.
    I, N = kernel_preds.shape
    kernel_preds = kernel_preds.view(I, N, 1, 1)
    seg_preds = F.conv2d(seg_preds, kernel_preds, stride=1).squeeze(0).sigmoid() # 得到 seg 3维  9ms
    #print("conv2d time {:.3f} ms".format((time.time() - start) * 1000))

    # mask.
    #seg_masks = seg_preds > mask_thr
    seg_masks = seg_preds > 0.5 # 大于阈值 # bool
    sum_masks = seg_masks.sum((1, 2)).float()

    # filter.
    keep = sum_masks > strides # 大于 seg的大小要大于strides
    if keep.sum() == 0:
        return None

    seg_masks = seg_masks[keep, ...] # bool
    seg_preds = seg_preds[keep, ...]
    sum_masks = sum_masks[keep]
    cate_scores = cate_scores[keep]
    cate_labels = cate_labels[keep]

    # mask scoring.
    seg_scores = (seg_preds * seg_masks.float()).sum((1, 2)) / sum_masks
    cate_scores *= seg_scores # 得分相乘 得到 置信度

    # sort and keep top nms_pre
    sort_inds = torch.argsort(cate_scores, descending=True) # 按得分高低进行排列
    if len(sort_inds) > max_per_img: # 取前100个
        sort_inds = sort_inds[:max_per_img]
    seg_masks = seg_masks[sort_inds, :, :]
    seg_preds = seg_preds[sort_inds, :, :]
    sum_masks = sum_masks[sort_inds]
    cate_scores = cate_scores[sort_inds]
    cate_labels = cate_labels[sort_inds]

    # Matrix NMS
    cate_scores = matrix_nms(seg_masks, cate_labels, cate_scores,
                                kernel='gaussian',sigma=2.0, sum_masks=sum_masks) #

    return seg_preds, cate_labels, cate_scores    


def matrix_nms(seg_masks, cate_labels, cate_scores, kernel='gaussian', sigma=2.0, sum_masks=None):
    """Matrix NMS for multi-class masks.

    Args:
        seg_masks (Tensor): shape (n, h, w)
        cate_labels (Tensor): shape (n), mask labels in descending order
        cate_scores (Tensor): shape (n), mask scores in descending order
        kernel (str):  'linear' or 'gauss' 
        sigma (float): std in gaussian method
        sum_masks (Tensor): The sum of seg_masks

    Returns:
        Tensor: cate_scores_update, tensors of shape (n)
    """
    n_samples = len(cate_labels)
    if n_samples == 0:
        return []
    if sum_masks is None:
        sum_masks = seg_masks.sum((1, 2)).float()
    seg_masks = seg_masks.reshape(n_samples, -1).float()
    # inter.
    inter_matrix = torch.mm(seg_masks, seg_masks.transpose(1, 0))
    # union.
    sum_masks_x = sum_masks.expand(n_samples, n_samples)
    # iou.
    iou_matrix = (inter_matrix / (sum_masks_x + sum_masks_x.transpose(1, 0) - inter_matrix)).triu(diagonal=1)
    # label_specific matrix.
    cate_labels_x = cate_labels.expand(n_samples, n_samples)
    label_matrix = (cate_labels_x == cate_labels_x.transpose(1, 0)).float().triu(diagonal=1)

    # IoU compensation
    compensate_iou, _ = (iou_matrix * label_matrix).max(0)
    compensate_iou = compensate_iou.expand(n_samples, n_samples).transpose(1, 0)

    # IoU decay 
    decay_iou = iou_matrix * label_matrix

    # matrix nms
    if kernel == 'gaussian':
        decay_matrix = torch.exp(-1 * sigma * (decay_iou ** 2))
        compensate_matrix = torch.exp(-1 * sigma * (compensate_iou ** 2))
        decay_coefficient, _ = (decay_matrix / compensate_matrix).min(0)
    elif kernel == 'linear':
        decay_matrix = (1-decay_iou)/(1-compensate_iou)
        decay_coefficient, _ = decay_matrix.min(0)
    else:
        raise NotImplementedError

    # update the score.
    cate_scores_update = cate_scores * decay_coefficient
    return cate_scores_update


def vis_seg(image_raw, result, score_thresh, output):

    img_show = image_raw # no pad
    seg_show = img_show.copy()

    ori_h,ori_w,_ = image_raw.shape


    if result!=None:
        seg_label = result[0].cpu().numpy() # seg
        output_scale = [ ori_w/seg_label.shape[2] , ori_h/seg_label.shape[1] ]
        #seg_label = seg_label.astype(np.uint8) # 变成int8
        cate_label = result[1] # cate_label
        cate_label = cate_label.cpu().numpy()
        score = result[2].cpu().numpy() # cate_scores
        
        vis_inds = score > score_thresh # 大于0.3
        seg_label = seg_label[vis_inds]
        num_mask = seg_label.shape[0]
        cate_label = cate_label[vis_inds]
        cate_score = score[vis_inds]

        for idx in range(num_mask):
            mask = seg_label[idx, :,:]
            # cur_mask = cv2.resize(cur_mask,(ori_w,ori_h))
            cur_mask = (mask> mask_thr).astype(np.uint8)
            if cur_mask.sum() == 0:
                continue
            mask_roi = cv2.boundingRect(cur_mask)

            draw_roi = (int(output_scale[0]*mask_roi[0]),int(output_scale[1]*mask_roi[1]),
                        int(output_scale[0]*mask_roi[2]),int(output_scale[1]*mask_roi[3]))
            
            now_mask = cv2.resize(mask[mask_roi[1]:mask_roi[1]+mask_roi[3],mask_roi[0]:mask_roi[0]+mask_roi[2]],(draw_roi[2],draw_roi[3]))
            now_mask = (now_mask> mask_thr).astype(np.uint8)
            color_mask = (np.random.randint(0,255),np.random.randint(0,255),np.random.randint(0,255))

            contours,_ = cv2.findContours(now_mask,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_NONE)

            draw_roi_mask = seg_show[ draw_roi[1]:draw_roi[1]+ draw_roi[3] , draw_roi[0]:draw_roi[0]+ draw_roi[2] ,:]
            cv2.drawContours(draw_roi_mask,contours,-1,color_mask,2)

            cur_cate = cate_label[idx]
            cur_score = cate_score[idx]

            label_text = class_names[cur_cate]

            vis_pos = (max(int(draw_roi[0]) - 10, 0), int(draw_roi[1])) #1ms
            #vis_pos = (max(int(center_x) - 10, 0), int(center_y)) #1ms
            cv2.rectangle(seg_show,(draw_roi[0],draw_roi[1]),(draw_roi[0]+ draw_roi[2],draw_roi[1]+ draw_roi[3]),(0,0,0),thickness=2)
            cv2.putText(seg_show, label_text, vis_pos,
                        cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 0))  # green 0.1ms

    cv2.imwrite(output,seg_show)



def main():

    args = argparse.ArgumentParser(description="trt pose predict")
    args.add_argument("--onnx_path",type=str)
    args.add_argument("--engine_path",type=str)
    args.add_argument("--image_path",type=str,default="demo/demo.jpg")
    
    args.add_argument("--mode",type=str,default="fp16")
    args.add_argument('--score_thr', type=float, default=0.3, help='score threshold for visualization')
    args.add_argument("--output",type=str,default="result.jpg")
    opt = args.parse_args()

    insize = [1,3,800,1344]
    model = TRT_model(insize,opt.onnx_path,opt.engine_path)
    preprocesser = Preprocessimage(insize)
    
    ############start inference##############
    image, image_raw,img_metas = preprocesser.process(opt.image_path)
    with torch.no_grad():
        result = model(image)
    vis_seg(image_raw, result, score_thresh=opt.score_thr, output=opt.output)
       




if __name__=="__main__":
    main()

测试效果

Model	GPU	Mode	Inference
R101	V100	fp16	35ms
R101	xavier	fp16	150ms

可见，经过fp16加速过的模型仍然很慢，毕竟输入图片的大小很大（1344*800），网络参数也不小，而且画出的后处理函数vis_seg耗时巨大（c++版的后处理可以将时间降低很多）。实际上距离工程应用仍然比较远，需要在模型结构优化下比较大的功夫。