快速搭建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两种包

六、总结

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

  • 20
    点赞
  • 21
    收藏
    觉得还不错? 一键收藏
  • 1
    评论
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值