笔记内容来自 B站 手写AI
一、用硬代码实现
GitHub - wang-xinyu/tensorrtx: Implementation of popular deep learning networks with TensorRT network definition API
安装python、cuda11.2、cudnn对应cuda11.2软件
1、在yolov5-master下训练完成后生成best.pt文件(训练时要清楚使用的什么类型训练的,可以是s\m\I\x)
2、在\tensorrtx-yolov5-v5.0\yolov5复制gen wts.py文件到\yolov5-master训练集下
3、使用anaconda工具运行该文件:python gen_wts.py -w yolov5l_best.pt
yolov5l_best.pt是训练生成的模型,通过该命令生成对应名字的wts,如yolov5l_best.wts
4、把生成的yolov5l_best.wts文件导入\tensorrtx-yolov5-v5.0\yolov5\build\Release文件下,使用指 令yolov5.exe -s yolov5l_best.wts yolov5l_best.engine l (通过yolov5.exe生成yolov5l best.engine,参数l根据训练模型来改,可以是s\m\I\x)
6、yolov5.exe -d yolov5l best.engine ./sample 测试生成的engine是否运行
win yolov5.7 tensorRT部署推理-CSDN博客
二、torch转TRT
为每个算子写转换器,反射Module.forward捕获输入输出和图结构。
三、ONNX
算子由官方提供,模型直接导出