Yolov5使用TensorRT加速(windows系统)

一、查看cuda版本

nvcc -V

二、下载安装包

        去英伟达官网下载cuda对应版本的TensorRT，先登录，安装过cuda时候应该大家就注册了nvidia帐号了。传送门

        下载完成后解压，进入文件夹打开终端，激活环境，执行以下命令，记得安装python对应的版本

conda activate yolov5

cd python

pip install tensorrt-8.6.1-cp38-none-win_amd64.whl
pip install tensorrt_dispatch-8.6.1-cp38-none-win_amd64.whl
pip install tensorrt_lean-8.6.1-cp38-none-win_amd64.whl

cd ../graphsurgeon
pip install graphsurgeon-0.4.6-py2.py3-none-any.whl

cd ../onnx_graphsurgeon
pip install onnx_graphsurgeon-0.3.12-py2.py3-none-any.whl

cd ../uff
pip install uff-0.6.9-py2.py3-none-any.whl

        找到cuda的环境目录，比如我的是C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.2

        将tensorrt文件夹的include目录下的所有文件复制到cuda的include下

        将lib目录下的所有文件复制到lib/x64

        将lib目录下的所有文件复制到bin

查看是否安装成功

        运行sample.py

cd samples/python/network_api_pytorch_mnist
python sample.py

        遇到提示：No module named 'cuda'

pip install cuda-python -i https://mirrors.aliyun.com/pypi/simple/

        这个函数是关于手写数字识别的，学ai的应该不陌生，如果运行后能训练及测试就成功了。

三、tensorrt推理

        转化模型，要经过漫长的等待......

python export.py --weights hat3.pt --data data/hat.yaml --include engine --device 0 --half

         运行

python detect.py --source 0 --weights hat3.engine

        成功！