【TensorRT】yolov5的pytorch模型转tensorrt模型（自己的数据集与模型）

骑猪看日落 2022-11-17 05:25 235阅读 0赞

代码参考：

[https://github.com/wang-xinyu/tensorrtx][https_github.com_wang-xinyu_tensorrtx]

[https://github.com/ultralytics/yolov5][https_github.com_ultralytics_yolov5]

**目录**

1.环境

2.安装TensorRT

3.pytorch转换TensorRT

4.TensorRT python接口验证

5.验证c++接口

--------------------

# 1.环境 #

yolov5环境配置参考：[https://blog.csdn.net/qq\_35975447/article/details/114061367][https_blog.csdn.net_qq_35975447_article_details_114061367]

# yolov5 environment
    
    ubuntu16.04
    cuda10.1
    cudnn7
    python3.6
     
     
    Cython
    matplotlib>=3.2.2
    numpy>=1.18.5
    opencv-python>=4.1.2
    Pillow
    PyYAML>=5.3
    scipy>=1.4.1
    tensorboard>=2.2
    torch>=1.7.0
    torchvision>=0.8.1
    tqdm>=4.41.0
    seaborn>=0.11.0
    easydict
    
    
    
    # tensorrt environment
    ubuntu16.04
    cuda10.1
    cudnn7
    python3.6
    
    tensorrt6.0.1.5
    
    # pip install 
    tensorrt
    pycuda

# 2.安装TensorRT #

安装参考：[https://github.com/wang-xinyu/tensorrtx/blob/master/tutorials/install.md][https_github.com_wang-xinyu_tensorrtx_blob_master_tutorials_install.md]

tensorrt安装参考：[https://zhuanlan.zhihu.com/p/64053177][https_zhuanlan.zhihu.com_p_64053177]，与我的：[https://blog.csdn.net/qq\_35975447/article/details/115632742][https_blog.csdn.net_qq_35975447_article_details_115632742]

# 3.pytorch转换TensorRT #

按照[https://github.com/wang-xinyu/tensorrtx/blob/master/yolov5/README.md][https_github.com_wang-xinyu_tensorrtx_blob_master_yolov5_README.md]中的过程，很容易成功。

1）转换 .wts文件

将/tensorrtx/yolov5/gen\_wts.py这个文件放到yolov5的根目录下，如果不修改文件中的路径，会读取weights/yolov5s.pt模型，在根目录下生成yolov5s.wts文件：

python gen_wts.py

2）.wts文件转.engine （tensorrt模型）

\----根据自己的类别修改yololayer.h：

# line 19
    
    static constexpr int CLASS_NUM = your_classes;

\----编译yolov5

cd yolov5
    mkdir build
    cd build
    cmake ..
    make

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM1OTc1NDQ3_size_16_color_FFFFFF_t_70][]

\----转换

将之前转好的模型复制到/tensorrtx/yolov5/目录下，然后运行下面的脚本：

./yolov5 -s ../yolov5s.wts ../yolov5s.engine s

![20210412190040333.png][]

# 4.TensorRT python接口验证 #

0）环境配置

下载对应的[TensorRT-6.0.1.5.Ubuntu-16.04.x86\_64-gnu.cuda-10.1.cudnn7.6.tar.gz][TensorRT-6.0.1.5.Ubuntu-16.04.x86_64-gnu.cuda-10.1.cudnn7.6.tar.gz]文件，解压后安装python目录下的whl文件。如：

pip install ./TensorRT-6.0.1.5/python/tensorrt-6.0.1.5-cp36-none-linux_x86_64.whl

安装pycuda：

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple pycuda --no-binary :all:

1）修改类别名称，改成自己的，这里可以将原来训练yolov5的目录下data/\*.yaml文件复制过来。

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM1OTc1NDQ3_size_16_color_FFFFFF_t_70 1][]

2）修改代码中的图片路径为自己的

# line 316
    input_image_paths = []
    
    # line 178
    image_raw = cv2.imread("./samples/" + input_image_path)

3）推理

将 2）中记录的图片放到samples目录下，将engine模型放到build目录下，然后运行命令进行推理：

python yolov5_trt.py

成功的话，会在yolov5目录下，生成以自己的文件命名+"\_output"的检测结果。我自己看了下，发现是可以的，应该没毛病。

# 5.验证c++接口 #

进入build目录下，运行：

./yolov5 -d yolov5s.engine ../samples

[https_github.com_wang-xinyu_tensorrtx]: https://github.com/wang-xinyu/tensorrtx
[https_github.com_ultralytics_yolov5]: https://github.com/ultralytics/yolov5
[https_blog.csdn.net_qq_35975447_article_details_114061367]: https://blog.csdn.net/qq_35975447/article/details/114061367
[https_github.com_wang-xinyu_tensorrtx_blob_master_tutorials_install.md]: https://github.com/wang-xinyu/tensorrtx/blob/master/tutorials/install.md
[https_zhuanlan.zhihu.com_p_64053177]: https://zhuanlan.zhihu.com/p/64053177
[https_blog.csdn.net_qq_35975447_article_details_115632742]: https://blog.csdn.net/qq_35975447/article/details/115632742
[https_github.com_wang-xinyu_tensorrtx_blob_master_yolov5_README.md]: https://github.com/wang-xinyu/tensorrtx/blob/master/yolov5/README.md
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM1OTc1NDQ3_size_16_color_FFFFFF_t_70]: /images/20221022/0f23f47f70fe46a5b3ccd1f3ebc76eec.png
[20210412190040333.png]: /images/20221022/997bee24dcca48c0a4c5a8a434944ccc.png
[TensorRT-6.0.1.5.Ubuntu-16.04.x86_64-gnu.cuda-10.1.cudnn7.6.tar.gz]: https://developer.nvidia.com/compute/machine-learning/tensorrt/secure/6.0/GA_6.0.1.5/tars/TensorRT-6.0.1.5.Ubuntu-16.04.x86_64-gnu.cuda-10.1.cudnn7.6.tar.gz
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzM1OTc1NDQ3_size_16_color_FFFFFF_t_70 1]: /images/20221022/8c4bbe48d79e478ca468515fe7bcd18c.png