yolov8实战第三天——yolov8TensorRT部署（python推理）（保姆教学）

军超

在上一篇中我们使用自己的数据集训练了一个yolov8检测模型，best.py。yolov8实战第一天——yolov8部署并训练自己的数据集（保姆式教程）-CSDN博客yolov8实战第二天——yolov8训练结果分析（保姆式解读）-CSDN博客接下要对best.py进行TensorRT优化并部署。TensorRT是一种高性能深度学习推理优化器和运行时加速库，可以为深度学习应用提供低延迟、高吞吐率的部署推理。TensorRT可用于对超大规模数据中心、嵌入式平台或自动驾驶平台进行推理加速。TensorRT现已能支持TensorFlow、Caffe、Mxnet、Pytorch等几乎所有的深度学习框架，将TensorRT和NVIDIA的GPU结合起来，能在几乎所有的框架中进行快速和高效的部署推理。一般的深度学习项目，训练时为了加快速度，会使用多GPU分布式训练。但在部署推理时，为了降低成本，往往使用单个GPU机器甚至嵌入式平台（比如NVIDIAJetson）进行部署，部署端也要有与训练时相同的深度学习环境，如caffe，TensorFlow等。由于训练的网络模型可能会很大（比如，inception，resnet等），参数很多，而且部署端的机器性能存在差异，就会导致推理速度慢，延迟高。这对于那些高实时性的应用场合是致命的，比如自动驾驶要求实时目标检测，目标追踪等。为了提高部署推理的速度，出现了很多模型优化的方法，如：模型压缩、剪枝、量化、知识蒸馏等，这些一般都是在训练阶段实现优化。而TensorRT 则是对训练好的模型进行优化，通过优化网络计算图提高模型效率。一、安装TensorRT  Login|NVIDIADeveloper下载TensorRT。我下载的是8.6里画黑线的那个。 将TensorRT-8.6.1.6\include中头文件copy到C:\ProgramFiles\NVIDIAGPUComputingToolkit\CUDA\v11.7\include将TensorRT-8.6.1.6\lib中所有lib文件copy到C:\ProgramFiles\NVIDIAGPUComputingToolkit\CUDA\v11.7\lib\x64将TensorRT-8.6.1.6\lib中所有dll文件copy到C:\ProgramFiles\NVIDIAGPUComputingToolkit\CUDA\v11.7\bin在python文件夹中找到适合自己的。pipinstalltensorrt-8.6.1-cp310-none-win_amd64.whl 至此TensorRT安装完成。 二、pt转onnx：GitHub-triple-Mu/YOLOv8-TensorRT:YOLOv8usingTensorRTaccelerate!参考着这个，下载，安装环境后。安装onnx：pipinstallonnx-ihttps://pypi.tuna.tsinghua.edu.cn/simplepipinstallonnxsim-ihttps://pypi.tuna.tsinghua.edu.cn/simplepipinstallonnxruntime-ihttps://pypi.tuna.tsinghua.edu.cn/simple生成onnx： pythonexport-det.py--weightsyolov8n.pt--iou-thres0.65--conf-thres0.25--topk100--opset11--sim--input-shape13640640--devicecuda:0使用上篇文章中的老鼠模型做了测试： onnx的测试代码：importonnxruntimeasrtimportnumpyasnpimportcv2importmatplotlib.pyplotaspltdefnms(pred,conf_thres,iou_thres):conf=pred[...,4]>conf_thresbox=pred[conf==True]cls_conf=box[...,5:]cls=[]foriinrange(len(cls_conf)):cls.append(int(np.argmax(cls_conf[i])))total_cls=list(set(cls))output_box=[]foriinrange(len(total_cls)):clss=total_cls[i]cls_box=[]forjinrange(len(cls)):ifcls[j]==clss:box[j][5]=clsscls_box.append(box[j][:6])cls_box=np.array(cls_box)box_conf=cls_box[...,4]box_conf_sort=np.argsort(box_conf)max_conf_box=cls_box[box_conf_sort[len(box_conf)-1]]output_box.append(max_conf_box)cls_box=np.delete(cls_box,0,0)whilelen(cls_box)>0:max_conf_box=output_box[len(output_box)-1]del_index=[]forjinrange(len(cls_box)):current_box=cls_box[j]interArea=getInter(max_conf_box,current_box)iou=getIou(max_conf_box,current_box,interArea)ifiou>iou_thres:del_index.append(j)cls_box=np.delete(cls_box,del_index,0)iflen(cls_box)>0utput_box.append(cls_box[0])cls_box=np.delete(cls_box,0,0)returnoutput_boxdefgetIou(box1,box2,inter_area):box1_area=box1[2]*box1[3]box2_area=box2[2]*box2[3]union=box1_area+box2_area-inter_areaiou=inter_area/unionreturnioudefgetInter(box1,box2):box1_x1,box1_y1,box1_x2,box1_y2=box1[0]-box1[2]/2,box1[1]-box1[3]/2,\box1[0]+box1[2]/2,box1[1]+box1[3]/2box2_x1,box2_y1,box2_x2,box2_y2=box2[0]-box2[2]/2,box2[1]-box1[3]/2,\box2[0]+box2[2]/2,box2[1]+box2[3]/2ifbox1_x1>box2_x2orbox1_x2box2_y2orbox1_y2None:device=torch.device(args.device)Engine=TRTModule(args.engine,device)H,W=Engine.inp_info[0].shape[-2:]#setdesiredoutputnamesorderEngine.set_desired(['num_dets','bboxes','scores','labels'])images=path_to_list(args.imgs)save_path=Path(args.out_dir)ifnotargs.showandnotsave_path.exists():save_path.mkdir(parents=True,exist_ok=True)forimageinimages:save_image=save_path/image.namebgr=cv2.imread(str(image))draw=bgr.copy()bgr,ratio,dwdh=letterbox(bgr,(W,H))rgb=cv2.cvtColor(bgr,cv2.COLOR_BGR2RGB)tensor=blob(rgb,return_seg=False)dwdh=torch.asarray(dwdh*2,dtype=torch.float32,device=device)tensor=torch.asarray(tensor,device=device)#inferencedata=Engine(tensor)bboxes,scores,labels=det_postprocess(data)ifbboxes.numel()==0:#ifnoboundingboxprint(f'{image}:noobject!')continuebboxes-=dwdhbboxes/=ratiofor(bbox,score,label)inzip(bboxes,scores,labels):bbox=bbox.round().int().tolist()cls_id=int(label)cls=CLASSES[cls_id]color=COLORS[cls]cv2.rectangle(draw,bbox[:2],bbox[2:],color,2)cv2.putText(draw,f'{cls}:{score:.3f}',(bbox[0],bbox[1]-2),cv2.FONT_HERSHEY_SIMPLEX,0.75,[225,255,255],thickness=2)ifargs.show:cv2.imshow('result',draw)cv2.waitKey(0)else:cv2.imwrite(str(save_image),draw)defparse_args()->argparse.Namespace:parser=argparse.ArgumentParser()parser.add_argument('--engine',type=str,help='Enginefile')parser.add_argument('--imgs',type=str,help='Imagesfile')parser.add_argument('--show',action='store_true',help='Showthedetectionresults')parser.add_argument('--out-dir',type=str,default='./output',help='Pathtooutputfile')parser.add_argument('--device',type=str,default='cuda:0',help='TensorRTinferdevice')args=parser.parse_args()returnargsif__name__=='__main__':args=parse_args()main(args)