快速搭建tensorRT环境及pt转为tensorrt推理【入门】

  pt权重转为tensorrt，在这个过程踩过很多坑，特别是在环境搭建这部分，要不就是系统有问题、编译出错、pycuda安装、gcc各种坑，所以本篇文章，是想让大家少走弯路，实现最快速度的模型转换。

一、什么是tensorrt

  TensorRT是英伟达针对自家平台做的一个加速包，是可以在NVIDIA各种GPU硬件平台下运行的一个C++推理框架。根据官方文档，使用TensorRT，在CPU或者GPU模式下其可提供10X乃至100X的加速。TensorRT 的限制就是只能运行在 NVIDIA 显卡上
TensorRT主要做了这么两件事情，来提升模型的运行速度：
（1）TensorRT支持INT8和FP16的计算。深度学习网络在训练时，通常使用 32 位或 16 位数据。TensorRT则在网络的推理时选用不这么高的精度，达到加速推断的目的
（2）TensorRT对于网络结构进行了重构，把一些能够合并的运算合并在了一起，针对GPU的特性做了优化

一句话总结：TensorRT是nvidia的加速包，能够在推理精度和网络结构合并重构这两个方面，加速模型推理。

二、如何搭建tensorrt环境

  TensorRT的版本与CUDA还有CUDNN版本是密切相关的，不是随便装一个进去就行。推荐在镜像内部安装tensorrt，以后模型转换在镜像内部转换，不影响其他环境。那如果有个镜像直接帮我装好了，我岂不是不用自己装！！！
  先查下自己设备的cuda、cudnn版本

查询cuda版本命令：
  nvcc -V
  nvidia-smi
  cat /usr/local/cuda/version.txt
查询cudnn版本（没返回则说明没有安装）
  cat /usr/local/cuda/include/cudnn.h | grep CUDNN_MAJOR -A 2

  在dockerhub查找tensorrt关键字，可以看到有一个hakuyyf/tensorrtx，里面有适配cuda11和cuda10的镜像，我的cuda是10.1的版本，拉取对应的镜像，docker pull hakuyyf/tensorrtx:trt7_cuda10

三、tensorrt的转换方法

  tensorrt的转换方式有三种，参考https://zhuanlan.zhihu.com/p/371239130
1、使用TF-TRT，将TensorRT集成在TensorFlow中
2、使用ONNX2TensorRT，即ONNX转换trt的工具
3、手动构造模型结构，然后手动将权重信息挪过去，非常灵活但是时间成本略高，有大佬已经尝试过了：tensorrtx

四、tensorrtx

  上述的使用方法2，pt先转成onnx，在转换为tensorrt，可以参考https://zhuanlan.zhihu.com/p/671320938，对C++不了解，看不懂。要是有现成的就好！那就是方法3，tensorrtx
  tensorrtx是wangxinyu大佬自己写的通过tensorrt的API直接构建网络的进行模型部署的方案，有各类模型的推理方法，例如yolo系列、dbnet、arcface等。
  这里需要注意下，如果是拉取的cuda11对应的tensorrt镜像，可以不做修改，如果是cuda10对应的tensorrt镜像，需要重新构建镜像，安装opencv。dockerfile可查看tensorrtx里面的docker/trt7目录

重新构建镜像的命令：

cd docker/trt7
docker build -t trt7_cuda10_opencv34 ./

启动镜像，并把tensorrtx的代码放到经镜像里面

# 启动
docker run -it --gpus all trt7_cuda10_opencv34
# 查看运行的容器ID
docker ps
# 将tensorrtx代码放入容器
docker cp tensorrtx 容器ID:/

五、案例：yolov5转换tensorrt

  首先生成wts格式的文件，再进行编译。
  wts格式的文件，第一个数字代表层，后面每一层都是这种格式[weight name] [value count = N] [value1] [value2], …, [valueN]

# wts生成：
# 下载yolov5的源码https://github.com/ultralytics/yolov5
# 下载权重yolov5s.pt
# 将tensorrtx源码里面的yolov5/gen_wts.py，复制到yolov5目录
cd tensorrtx/yolov5
cp gen_wts.py yolov5源码目录
# wts生成
cd yolov5源码目录
python gen_wts.py -w yolov5s.pt -o yolov5s.wts
# 可以看到yolov5的目录下有yolov5s.wts文件
# 将yolov5s.wts放入到镜像中进行编译
cd yolov5源码目录
docker cp yolov5s.wts 容器ID:/tensorrtx/yolov5
# 进入容器的yolov5目录
cd /tensorrtx/yolov5
# 创建build目录
mkdir build
cd build
# 将yolov5s.wts放入到build目录下
cp ../yolov5s.wts .
# 编译
cmake ..
make
# 序列化生成.engine文件
./yolov5_det -s yolov5s.wts yolov5s.engine s
# 推理
./yolov5_det -d yolov5s.engine ../images
# 可以看到推理的图像

  作者也提供了用python如何推理tensorrt权重，可以看下tensorrtx/yolov5yolov5_det_cuda_python.py、yolov5_det_trt.py，分别用了cuda和pycuda两种包

六、总结

  如果文章对您有所帮助，记得点赞、收藏、评论探讨✌️