仅用于站内搜索，没有排版格式，具体信息请跳转上方微信公众号内链接
以下文章来源于知乎
作者：启思
链接：https ://zhuanlan. zhihu.com/p/527238167
本文仅用于学术分享，如有侵权，请联系后台作删文处理
导读
模型部署作为人工智能的落地的“最后一步”，也是算法能够转换为生产力的重要环节。因此本文分享TensorRT的部署流程，希望能对各位读者有所帮助。
前段时间用TensorRT部署了一套模型，速度相比Python实现的版本快了20多倍，中间踩了许多坑，但是最后发现流程其实相当简单，特此记录一下踩坑过程。
01
TensorRT
这东西就是NVidia在自家显卡上做了一个深度学习inference加速的框架，只要你把训练好的模型参数和结构告诉他，他就能自动帮你优化（硬件相关），以达到最快速度。
这涉及两个问题：
应该以什么模型格式把模型喂给TensorRT？
如何使用TensorRT优化后的模型？
对于第一个问题：现在的深度学习框架非常多，不止常用的pytorch/tensorflow，而即使是同一种框架还可以使用不同的编程语言实现。让TensorRT对每一个框架都直接支持，显然是不可能的。
TensorRT只需要知道网络的结构和参数即可，它支持三种转换入口：
TF-TRT，要求是TensorFlow模型
ONNX模型格式
使用TensorRTAPI手动把模型搭起来，然后把参数加载进去
第一种不够灵活，第三种比较麻烦，所以最省事方便的就是第二种方法。本文介绍第二种。
ONNX就是一个通用的神经网络格式，一个. onnx文件内包含了网络的结构和参数。甭管是用什么深度学习框架写的网络，只要把模型导出成ONNX格式，就跟原本的代码没有关系了。
转成ONNX格式还没有被优化，需要再使用TensorRT读取它并优化成TensorRTEngine。优化参数也在这一步指定。
对于第二个问题：得到的TensorRTEngine是硬件相关的，之后跑模型只需要运行这个Engine即可。调用TensorRTEngine需要使用TensorRTRuntimeAPI。
所以整个逻辑就是：
把你的模型导出成ONNX格式。
把ONNX格式模型输入给TensorRT，并指定优化参数。
使用TensorRT优化得到TensorRTEngine。
使用TensorRTEngine进行inference。
02
你需要做的
把模型导出成ONNX格式。
安装TensorRT和CUDA。注意二者和driver的版本号对应，我用的是ZIP安装，跟着这个把流程走一遍。
https ://docs. nvidia.com/deeplearning/tensorrt/install-guide/index. html#installing-zip
设置优化参数，使用TensorRT把ONNX优化成Engine，得到当前硬件上优化后的模型。
使用TensorRTRuntimeAPI进行inference。
官方文档，写的很详细了。https ://docs. nvidia.com/deeplearning/tensorrt/developer-guide/#overview
03
ONNX转换
Pytorch自带导出方法torch. onnx.export。
TensorFlow推荐使用这个https ://github. com/onnx/tensorflow-onnx，从checkpoint（.meta后缀）导出只需要
即可，tf1. x/2. x均可以用。
这里一定要注意INPUT_NAME和OUTPUT_NAME有没有写对，它决定了网络的入口和出口，转换ONNX错误/导出Engine错误很有可能是这个没指定对。
不确定的话可以用onnxruntime看一看对不对
还可以在这个网站可视化导出的ONNX：https ://netron. app/
04
使用trtexec. exe测试
参考：
TensorRT安装流程走完之后就能在TensorRT-x-x-x-x/bin/文件夹下看到trtexec. exe。
这个命令：
./trtexec–onnx=model. onnx
把onnx模型优化成Engine，然后多次inference后统计并报时。
当然可以把Engine文件导出，使用–saveEngine参数
./trtexec–onnx=model. onnx–saveEngine=xxx. trt
一般来说模型的第一维大小是任意的（batchsize维度），而TensorRT不能把任意batchsize都加速。可以指定一个输入范围，并且重点优化其中一个batchsize。例如网络输入格式是[-1,3,244,244]：
这个input对应于INPUT_NAME:0。
还可以降低精度优化速度。一般来说大家写的模型内都是float32的运算，TensorRT会默认开启TF32数据格式，它是截短版本的FP32，只有19bit，保持了fp16的精度和fp32的指数范围。
另外，TensorRT可以额外指定精度，把模型内的计算转换成float16或者int8的类型，可以只开一个也可以两个都开，trtexec会倾向于速度最快的方式（有些网络模块不支持int8）
./trtexec–onnx=model. onnx–saveEngine=xxx. trt–int8–fp16
trtexec还提供了–best参数，这相当于–int8–fp16同时开。
一般来说，只开fp16可以把速度提一倍并且几乎不损失精度；但是开–int8会大大损失精度，速度会比fp16快，但不一定能快一倍。
int8优化涉及模型量化，需要校准（calibrate）提升精度。TensorRT有两种量化方法：训练后量化和训练中量化。二者的校准方法不同，精度也不同，后者更高一些。具体参考
https ://docs. nvidia.com/deeplearning/tensorrt/developerguide/index. html#working-with-int8
trtexec采用的是训练后量化，写起来更方便一些。不过看看源码就能发现，因为trtexec只为了测试速度，所以校准就象征性做了一下，真想自己部署int8模型还得自己写校准。
05
使用TensorRTC++API
trtexec只能看模型最快能跑多快，它是不管精度的，如果真想实际部署上又快又好的模型还是要自己调TensorRT的API。
可以用C++API、PythonAPI、TF-TRTRuntime，因为TF-TRT有局限性，C++API的速度比PythonAPI快，所以我选择C++API。三者区别可以参考：
参考TensorRT的sample自己写并不难。把ONNX转换成TensorRTEngine的代码是：
如果要加上fp16或者int8优化，需要在serialized之前，修改config
这里我把输入的范围规定成唯一了。int8这里的校准是我仿照trtexec写的，具体看下一节。
如果直接从文件中读取Engine就更简单了
把Engine保存到文件：
使用时，需要告诉网络输入输出在显存上的cuda指针。
这是同步版的，当然还有异步版的enqueueV2。
06
使用TensorRTC++API实现int8校准
这里用的还是训练后校准。逻辑是：搞一些真实输入数据（不需要输出），告诉TensorRT，它会根据真实输入数据的分布来调整量化的缩放幅度，以最大程度保证精度合适。
理论上校准数据越多，精度越高，但实际上不需要太多数据，TensorRT官方说500张图像就足以校准ImageNet分类网络。
我的校准部分是
就是搞了4096个batch的数据（这个数看着设就行），后面那个类是自己实现的，负责告诉校准器每个batch的数据是什么
继承的类是IInt8EntropyCalibrator2，这个得根据需要选择，不同类型的网络不一样，详见
https ://docs. nvidia.com/deeplearning/tensorrt/developerguide/index. html#enable_int8_c
这里面有一个read和write函数我没实现，它们负责从文件中读和写calibrationcache。如果用这个cache那么省去了生成数据和网络跑数据的时间，生成Engine时会更快一些。
欢迎加入《AI未来星球》，一起成长
扫描下方二维码即可加入~
真诚分享AI落地过程（AI商机->项目签约->算法开发->产品开发->实施运维）中的各方面经验和踩过的坑。
你可以获得什么？
1、大白之前花费10W+购买，AI行业各场景私有数据集下载，星球内倾情分享；2、AI行业研发、产品、商业落地问题咨询（目前AI公司创业中），都可获高质量解答，有效期一年，无限次提问，有问必答。3、定期邀请AI行业各类嘉宾分享，创业/商业等方面的经验！
帮助你解决遇到的实际问题，升职加薪！
大家一起加油！
