最近考虑对已经训练好的Pytorch模型进行Float16优化,以达到部署时inference减小耗时,显存开销。可是很多博客写的乱七八糟,我这里整理一份自己安装的过程。供大家参考。
首先我们需要查看自己服务器或者GPU是什么版本的驱动,然后才可以安装对应版本的Tensorrt
[~]$ 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版本。
pip install 'pycuda>=2017.1.1'
解压完成之后会有两个文件夹如下:
并把lib添加至环境变量
LD_LIBRARY_PATH=~/TensorRT-6.0.1.5/lib:$LD_LIBRARY_PATH
安装 Tensorrt
安装 uff
安装 graphsurgeon
测试安装是否成功
>>> import tensorrt
>>> tensorrt.__version__
'6.0.1.5'
您可以训练模型或从一个框架中加载预先训练的模型,将此模型导出为 ONNX,然后将此模型导入到另一个框架中
libtorch
,但使用libtorch推理的前提是:首先trace已经训练好的模型。这个代价实际上比较高,每个模型都需要做特定的trace,很多时候,如果有if-else分支,还得分离模型,包括tensor返回值都需要特定的模式。目前pytorch比较成熟的方式:①trace 模型,然后 libtorch 前向 ② pytorch模型转到onnx模型,然后使用tensorrt部署;第二种部署方式更加有效率
训练好的网络部署需要使用网络解析器:有三种
ONNX parser, Caffe parser, or TensorFlow/UFF parser.
directly from Caffe, or from other frameworks via the UFF or ONNX formats
.