使用tensorrt 加速Ultra-Fast-Lane-Detection项目

使用tensorrt 加速Ultra-Fast-Lane-Detection项目

运行环境：

项目	主机版本(on ubuntu 18.04)	AGX版本(Jetson 4.3)
cuda 10.0	10.0.130	-
cudnn7	7.6.5.32	-
tensorrt	7.0.0.11-1+cuda10.0	-
pytorch	1.4.0+cu100	1.2.0
torchvision	0.5.0+cu100	0.4.0

项目地址：
https://gitee.com/n__n/Ultra-Fast-Lane-Detection.git

pth 转 onnx

#!/usr/bin/python3

import torch
import torch.onnx as tox
from torch.autograd import Variable
from model.model import parsingNet
from utils.common import merge_config
from utils.dist_utils import dist_print

if __name__ == "__main__":
    torch.backends.cudnn.benchmark = True
    args, cfg = merge_config()

    dist_print('start testing...')
    assert cfg.backbone in ['18', '34', '50', '101',
                            '152', '50next', '101next', '50wide', '101wide']

    if cfg.dataset == 'CULane':
        cls_num_per_lane = 18
    else:
        raise NotImplementedError

    net = parsingNet(pretrained=False, backbone=cfg.backbone, cls_dim=(
        cfg.griding_num + 1, cls_num_per_lane, 4), use_aux=False).cuda()

    state_dict = torch.load(cfg.test_model, map_location='cuda')['model']
    compatible_state_dict = {
   }
    for k, v in state_dict.items():
        if 'module.' in k