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()
01-30
865
![](https://csdnimg.cn/release/blogv2/dist/pc/img/readCountWhite.png)
07-13
1739
![](https://csdnimg.cn/release/blogv2/dist/pc/img/readCountWhite.png)
07-25
333
![](https://csdnimg.cn/release/blogv2/dist/pc/img/readCountWhite.png)
09-17