当前位置: 首页 > 工具软件 > TensorRT > 使用案例 >

TensorRT 安装与使用及Pytorch的TensorRT部署

景承教
2023-12-01

1. Tensorrt安装

最近考虑对已经训练好的Pytorch模型进行Float16优化,以达到部署时inference减小耗时,显存开销。可是很多博客写的乱七八糟,我这里整理一份自己安装的过程。供大家参考。

首先我们需要查看自己服务器或者GPU是什么版本的驱动,然后才可以安装对应版本的Tensorrt

1.1 查看cuda版本

[~]$ nvcc -V

nvcc: NVIDIA ® Cuda compiler driver
Copyright © 2005-2017 NVIDIA Corporation
Built on Fri_Sep__1_21:08:03_CDT_2017
Cuda compilation tools, release 9.0, V9.0.176

下载 Tensorrt6.x
注:下载时,需要注册Nvidia的developer的账号,填写对应的开发者调查问卷信息。不一定非要下载6.x版本,目前已更新到7.x,不过考虑环境兼容问题,所以先选择6.x版本。

1.2 安装

pip install 'pycuda>=2017.1.1'

1.3

解压完成之后会有两个文件夹如下:

  • TensorRT-6.0.1.5
  • TensorRT-6.0.1.5.CentOS-7.6.x86_64-gnu.cuda-9.0.cudnn7.6.tar.gz

并把lib添加至环境变量
LD_LIBRARY_PATH=~/TensorRT-6.0.1.5/lib:$LD_LIBRARY_PATH

安装 Tensorrt

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

安装 uff

  • cd TensorRT-6.0.1.5/uff
  • pip install uff-0.6.5-py2.py3-none-any.whl

安装 graphsurgeon

  • cd TensorRT-6.0.1.5/graphsurgeon
  • pip install graphsurgeon-0.4.1-py2.py3-none-any.whl

测试安装是否成功

>>> import tensorrt
>>> tensorrt.__version__
'6.0.1.5'

2. 什么是onnx

您可以训练模型或从一个框架中加载预先训练的模型,将此模型导出为 ONNX,然后将此模型导入到另一个框架中

  • 个人理解:onnx是做各种框架下的模型转换
  • 例如:我用pytorch训练了一个model,现在有人想用这个做finetune;不过,他熟悉的tensorflow
  • 路径可以是: torch.pt -> model.onnx -> flow.pkl 然后就可以用这个model了。

2.1 pytorch与onnx、tensorRT什么关系

  • pytorch是训练框架,与caffe、tensorflow、Mxnet等一样
  • 很多时候,我们部署pytorch的model使用libtorch,但使用libtorch推理的前提是:首先trace已经训练好的模型。这个代价实际上比较高,每个模型都需要做特定的trace,很多时候,如果有if-else分支,还得分离模型,包括tensor返回值都需要特定的模式。
  • onnx可以认为是一个文件协议(表示机器学习模型的开放标准), 许多框架(包括 TensorFlow、PyTorch、SciKit-Learn、Keras、Chainer、MXNet 和 MATLAB)中的模型都可以导出或转换为标准 ONNX 格式。 模型采用 ONNX 格式后,可在各种平台和设备上运行。可以支持不同平台的CPU,GPU推理
  • tensortRT 是一个推理优化器,是一个C++库

目前pytorch比较成熟的方式:①trace 模型,然后 libtorch 前向 ② pytorch模型转到onnx模型,然后使用tensorrt部署;第二种部署方式更加有效率

3. Pytorch model 的Tensorrt部署

3.1 python环境 部署

3.2 C++环境 部署

训练好的网络部署需要使用网络解析器:有三种

  • Generally, this means you need to take the saved neural network and parse it from its saved format into TensorRT using the ONNX parser, Caffe parser, or TensorFlow/UFF parser.
  • Networks can be imported directly from Caffe, or from other frameworks via the UFF or ONNX formats.

4. Caffe 模型的Tensorrt部署

 类似资料: