1.下载必要的脚本
git clone https://github.com/wang-xinyu/tensorrtx.git git clone -b archive https://github.com/ultralytics/yolov3.git 将以上两个脚本放置同一目录下; 或者采用笔者上传整理好的脚本文件 链接地址:https://download.csdn.net/download/weixin_44533869/21026929
2.配置环境
- 下载pytorch离线安装包,放置pkg目录下。 torch-1.6.0-cp36-cp36m-linux_aarch64.whl
- 下载torchvision源码库,放置pkg目录下。
git clone --branch v0.7.0 https://github.com/pytorch/vision torchvision - 执行安装脚本
bash install.sh
3.Yolo转TensorRT Engine流程
yolo -> wts(Pytorch Yolo) -> engine(TensorRT engine)
4. yolo转wts
进入yolov2wts目录
- 将.cfg拷贝至cfg目录,将.weights拷贝至weights目录。(两个文件前缀必须相同,如yolov4_final.cfg,yolov4_final.weight,以下也以此为例)。
- 执行生成yolov4_safe.wts命令。
python3 gen_wts.py yolov4_final
此处可以在本地执行,也可以在远端服务器执行
- wts目录下将生成yolov4_safe.wts后缀的文件。
5. 下载Tensorrt
https://developer.nvidia.com/nvidia-tensorrt-7x-download 此处以Tensorrt-7x为例 笔者采用TAR格式的package包
6.安装Tensorrt
- 解压TensorRT-7.2.3.4.Ubuntu-18.04.x86_64-gnu.cuda-11.1.cudnn8.1.tar.gz
- 将下面环境变量写入环境变量文件~/.bashrc并保存
vim ~/.bashrc
export LD_LIBRARY_PATH=TensorRT解压路径/lib:$LD_LIBRARY_PATH
例如:
export LD_LIBRARY_PATH=/home/zhougen001/TensorRT-7.2.3.4/lib:$LD_LIBRARY_PATH
- 使刚刚修改的环境变量文件生效
source ~/.bashrc
- 复制TensorRT路径下/lib、/include文件夹到对应系统文件夹
sudo cp -r ./lib
- 验证Tensorrt安装情况
如果要使用python接口的tensorrt,则需要安装pycuda
pip install pycuda
测试Tensorrt
$ python
>>>import tensorrt
>>>tensorrt.__version__
7. wts转TensorRT engine
进入wts2engine/yolov4目录
- 新建build目录,并进入build目录。
- 拷贝yolov4_safe.wts文件至build目录。
cp ../../../yolov2wts/wts/yolov4_safe.wts ./
- 修改yololayer.h头文件
static constexpr int CLASS_NUM = 12;
cmake ..
make
生成TensorRT engine文件
sudo ./yolov4 -s yolov4_safe
这里会生成yolov4_safe.engine文件。
- 验证模型输出
sudo ./yolov4 -d yolov4_safe ../../../sample/
在build目录下生成推理结果
|