深度学习模型服务端部署：Triton server+TensorRT【PyTorch格式----＞ONNX格式----＞TensorRT格式----＞利用Triton server部署】

左手的ㄟ右手 2023-10-14 11:23 40阅读 0赞

Tritonserver+TensorRT实现服务端部署

> 写在前面
> 
> 项目需要，因此要完成模型训练到实际服务端部署的全过程，但是之前只是专注于模型的训练，未接触过实战，就借此机会将训练好的模型部署全过程做一个记录

### 工具和环境需求[\#][Link 1] ###

我本地的电脑环境如下：

Python 3.8

PyTorch 1.12.1

GPU RTX 3060 Laptop

CUDA 11.2

`Tritonserver` 和 `TensorRT` 都是 Docker 拉取的镜像，使用的命令如下表：

docker pull nvcr.io/nvidia/tritonserver:21.08-py3
    docker pull nvcr.io/nvidia/tensorrt:21.08-py3

细心的应该已经注意到了，我拉取的`Tritonserver` 和 `TensorRT` 后缀都是一样的`21.08`，这是因为`Tritonserver`对能够使用的tensorrt文件是有要求的，如果运气不好，版本不对应有可能会导致`Tritonserver`运行不起来，为了避免踩坑，干脆就直接拉同样的即可。

如果嫌弃Docker的文件太大，也可以不拉取TensorRT的镜像，我给出下面的链接，可以自己对照表格去下载对应的版本。

https://docs.nvidia.com/deeplearning/triton-inference-server/release-notes/rel_22-07.html#rel_22-07

#### 工具介绍[\#][Link 2] ####

本来是想着直接用简单的Flask+TensorRT去做服务端的，但是奈何我比较懒，现在已经沉溺于Python，再去花费时间学一个Flask我估摸着时间不够，于是百度之后发现了Tritonserver，这个是NVIDIA自己做的开源产品，支持GRPC和HTTP，还说支持C API直接访问（这部分不了解，相关视频后续在官方教程里会给出）

最关键的是，官方在B站给了一个中文的教程，实战+框架理解，蛮详细的，下面是链接：

https://space.bilibili.com/1320140761/channel/collectiondetail?sid=493256

最后说一下为什么要选择TensorRT

因为如果要实际应用就要考虑模型的吞吐量，TensorRT对模型的加速是有很大的优化的，而且PyTorch到ONNX，再到TensorRT网上基本都有完善的教程和相关问题的解决方案。当前在我的笔记本上使用跑的还是蛮快的，我跑的是Video Swin-Transformer。

### 开搞[\#][Link 3] ###

#### PyTorch转ONNX[\#][Link 4] ####

我的模型是[Video Swin-Transformer][]，模型代码是基于[mmaction2][]框架做的。mmaction2是open-mmlab的项目之一，比较有名的应该就是mmdetection吧，我之前用过mmpose，数据处理模块用的时候真的爽的不要不要的，但是源码看起来也是绕迷宫一样的，多看几遍才能理清整个框架的运作流程。

然而官方的实现里没有给出Pytorch转ONNX的代码，只适用于其它模型，只好自己动手了咯。这部分并不难，给`Recognizer3D.py`文件添加一个`forward`函数，简化一下推理流程即可，就不详细描述了。

> P.S. 如果你也要导出Video Swin-Transformer的ONNX模型，不建议在本地上跑，当然，如果你的内存足够大20G以上，当我没说，这东西就是内存怪物。

#### ONNX优化——ONNXSIM[\#][Link 5] ####

ONNXSIM对ONNX做了进一步的改进，想看模型在优化结构的前后对比可以用这个网站——[netron.app][]，很方便。Video Swin-Transformer应用ONNXSIM之后效果出奇的好，原先的网络结构有点纷乱，优化后都快变成一条直线了。

给你们看一下效果图，非常直观，不，非常夸张

优化前：

[![17e42615c7045863b127da6045eb065d.png][]][17e42615c7045863b127da6045eb065d.png 1]

优化后：

[![d4944d0f8f91ee065ecacedc63420b17.png][]][d4944d0f8f91ee065ecacedc63420b17.png 1]

优化的代码很简单

python -m onnxsim input.onnx output.onnx

如果你有精力，想要进一步优化，可以选择去手动修改ONNX的结构，这个就不属于本文讨论的范围了。

> P.S. 我在Linux的GPU服务器上做ONNXSIM的时候，总是会显示Simplifying然后卡死，很奇怪，在自己本地就没这种事情

#### ONNX转TensorRT[\#][Link 6] ####

到了这一步，先别急着上手操作，先检查一下你的ONNX模型中涉及的算子有哪些，我推荐使用`onnxexplorer`

pip install onnxexplorer

然后用它查看onnx模型需要哪些算子

onnxexp glance -m input.onnx

[![a0a8ac8cd6c7271cbfdf84fd9185c4e5.png][]][a0a8ac8cd6c7271cbfdf84fd9185c4e5.png 1]

需要注意的是，上面的模型我是经过处理的，在处理前，是有一个算子叫做Mod，也就是取余操作，这个操作在TensorRT中是不支持的，你可以通过自己重新写一个算子，在TensorRT中叫Plugin，来解决这个问题，但是奈何我比较懒，我选择直接从根源解决问题，在Video Swin-Transformer源码中找到%这个运算符，用下面的函数代替：

def mod(x,y):
    	res = x - (x // y)*y
    	return res

至于如何查看ONNX转TensorRT支持的算子有哪些，请看下面的链接，Y是支持，N是不支持

https://github.com/onnx/onnx-tensorrt/blob/main/docs/operators.md

上述准备工作完成之后，即可上手转换，将ONNX文件转换为plan文件，这个因为Tritonserver似乎是只认这个`.plan`后缀，没试过其它后缀行不行。

trtexec --onnx=input.onnx --saveEngine=output.plan --buildOnly

这个`trtexec`工具是一个测试用的工具，非常好用，TensorRT自带的，去下载好的TensorRT文件夹下面的`bin`目录下就能找到。后面的`--buildOnly`选项表示跳过测试阶段，直接输出模型，如果你们想看自己模型的吞吐量，时延等等参数可以选择不加这个选项。  
最后补充一点，TensorRT似乎是不支持F.pad()这个张量操作的，会导致转TensorRT失败，官方的论坛里说用另一个工具做一个常量折叠就可以解决问题了，相关代码在下面

polygraphy surgeon sanitize --fold-constants input.onnx -o folded.onnx

#### Tritonserver[\#][Link 7] ####

到这里已经是最后一步了，但是仍然需要小心一点，就比如下面的命令

docker run --gpus all -it --rm --name triton_server \
     -p8001:8001 -p8002:8002 -p8003:8003 \
     -v D:\文档\dockerRepo:/workspace/dockerRepo \
     nvcr.io/nvidia/tritonserver:21.08-py3

仔细看官方的教学视频可以发现，他有一个多余的选项`--net host`，但是这个选项会导致后面的`-p`失效，你可以尝试一下，然后就会发现错误提醒。在我自己的电脑上8000端口被占用了，于是我就只好搞一个重映射，最好前后的端口号保持一致，这样子方便后续设置，也方便记忆。

`-v`选项是将本地仓库映射到docker中去，Repository的结构也有所要求，由于使用了TensorRT格式的模型，这一部分可以简便很多，可以不用写配置文件。大致的仓库架构如下：

-dockerRepo
    --model_1
    ---1
    ----model_1.plan
    ---model_1_labels.txt
    --model_2

其中数字代表着版本号，这是约定俗成的，也方便server做版本管理，然后运行下面的代码即可运行服务器了

tritonserver --model-repository=./dockerRepo --strict-model-config=false --http-port 8001 --grpc-port 8002 --metrics-port 8003

如果运行起来显示Model那一栏是空的，很有可能是TensorRT的版本和Tritonserver的版本不兼容，请仔细检查。

为了快速检验是否成功，可以使用curl去检查一下

curl -v localhost:8001/v2/health/ready

成功返回的话就是正常运行了，接下来编写client端即可，client的编写官方已经给出样例，自己学习即可，仔细展开篇幅太长不合适。

https://github.com/triton-inference-server

以上就是完整搭建推理服务器的全过程了，虽然简单，但是基本涵盖了大部分内容，进阶的部分请参考官方文档和官方教程，Triton server有很多有用的选项可以设置，比如batch的合成，队列的等待时间等等。

--------------------

[深度学习部署架构：以 Triton Inference Server（TensorRT）为例-腾讯云开发者社区-腾讯云][Triton Inference Server_TensorRT_-_-]

[https://www.cnblogs.com/seaman1900/p/16629649.html][https_www.cnblogs.com_seaman1900_p_16629649.html]

[NVIDIA英伟达的个人空间-NVIDIA英伟达个人主页-哔哩哔哩视频][NVIDIA_-NVIDIA_-]

[什么是Triton推理服务器？推理服务器的特点有哪些？ - 哔哩哔哩][Triton_ -]

[Link 1]: https://www.cnblogs.com/seaman1900/p/16629649.html#%E5%B7%A5%E5%85%B7%E5%92%8C%E7%8E%AF%E5%A2%83%E9%9C%80%E6%B1%82
[Link 2]: https://www.cnblogs.com/seaman1900/p/16629649.html#%E5%B7%A5%E5%85%B7%E4%BB%8B%E7%BB%8D
[Link 3]: https://www.cnblogs.com/seaman1900/p/16629649.html#%E5%BC%80%E6%90%9E
[Link 4]: https://www.cnblogs.com/seaman1900/p/16629649.html#pytorch%E8%BD%AConnx
[Video Swin-Transformer]: https://github.com/SwinTransformer/Video-Swin-Transformer.git
[mmaction2]: https://github.com/open-mmlab/mmaction2.git
[Link 5]: https://www.cnblogs.com/seaman1900/p/16629649.html#onnx%E4%BC%98%E5%8C%96onnxsim
[netron.app]: https://netron.app/
[17e42615c7045863b127da6045eb065d.png]: https://img-blog.csdnimg.cn/img_convert/17e42615c7045863b127da6045eb065d.png
[17e42615c7045863b127da6045eb065d.png 1]: https://img2022.cnblogs.com/blog/2138682/202208/2138682-20220827001904752-1459745600.png
[d4944d0f8f91ee065ecacedc63420b17.png]: https://img-blog.csdnimg.cn/img_convert/d4944d0f8f91ee065ecacedc63420b17.png
[d4944d0f8f91ee065ecacedc63420b17.png 1]: https://img2022.cnblogs.com/blog/2138682/202208/2138682-20220827001927177-243952328.png
[Link 6]: https://www.cnblogs.com/seaman1900/p/16629649.html#onnx%E8%BD%ACtensorrt
[a0a8ac8cd6c7271cbfdf84fd9185c4e5.png]: https://img-blog.csdnimg.cn/img_convert/a0a8ac8cd6c7271cbfdf84fd9185c4e5.png
[a0a8ac8cd6c7271cbfdf84fd9185c4e5.png 1]: https://img2022.cnblogs.com/blog/2138682/202208/2138682-20220827001943171-1550017261.png
[Link 7]: https://www.cnblogs.com/seaman1900/p/16629649.html#tritonserver
[Triton Inference Server_TensorRT_-_-]: https://cloud.tencent.com/developer/article/2126650
[https_www.cnblogs.com_seaman1900_p_16629649.html]: https://www.cnblogs.com/seaman1900/p/16629649.html
[NVIDIA_-NVIDIA_-]: https://space.bilibili.com/1320140761/channel/collectiondetail?sid=493256
[Triton_ -]: https://www.bilibili.com/read/cv14726032