笔记内容来自 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是否运行

二、torch转TRT

为每个算子写转换器，反射Module.forward捕获输入输出和图结构。

 

 三、ONNX

算子由官方提供，模型直接导出