Yolov5lite 转 trt

import sys 

sys.path.append("..") 

# -*- coding: utf-8 -* 

#################### 

from TLoger import * 

from TUtils import * 

#######Public######## 

import asyncio 

import math 

import numpy as np 

import cv2 

import os 

import time 

import json 

from collections import deque 

  

os.environ['CUDA_MODULE_LOADING'] = 'LAZY' 

os.environ['CUDA_LAUNCH_BLOCKING'] = '1' 

os.environ['CUDA_VISIBLE_DEVICES']="0,1,2,3" 

  

######TensorRT####### 

import torchvision 

import torchnvjpeg 

import torch 

import numpy as np 

import tensorrt as trt 

import cv2 

from torch2trt import TRTModule 

from torch2trt import torch2trt 

from copy import copy 

import pycuda.driver as cudadvr 

  

  

import models 

from models.experimental import attempt_load 

  

  

input="../weights/test.pt" 

output="../weights/test.trt" 

cuda=0 

batch=10 

depth=3 

shape=640 

  

try: 

    folder=os.path.dirname(output) 

    if not os.path.exists(folder): 

        os.mkdir(folder) 

  

    #需要导入 \models  和  \utils 两个文件夹 

    with torch.no_grad(): 

        #设置GPU卡 

        torch.cuda.set_device(cuda) 

        print("当前CUDA：",torch.cuda.current_device()) 

        model=attempt_load(input).eval().cuda() 

        x = torch.zeros((batch, depth, shape, shape)).cuda() 

        net = torch2trt(model, [x]) 

        torch.save(net.state_dict(), output) 

         

        print("保存路径：",output) 

except: 

    TLoger.err()