腾讯联合英伟达开源 TensorRT 插件自动生成工具 TPAT

来源: 投稿
作者: 罗奇奇
2022-03-25 16:47:00

2022 年 3 月 25 日,腾讯联合英伟达开发的 TensorRT 插件自动生成工具 TPAT 正式宣布开源。

TensorRT 是当前应用最广的 GPU 推理框架,但由于支持的算子数量有限,用户面临手写插件以支持算子的痛点。TPAT 能够支持开放神经网络交换 (ONNX) 格式所有的算子,端到端生成 TensorRT 插件,在解放人力成本的同时,性能对比手写毫不逊色。

TPAT Github 地址:https://github.com/Tencent/TPAT

背景

TensorRT 是当今最快的 GPU 推理引擎,可以让深度学习模型在 GPU 上实现低延迟、高吞吐量的部署,支持 Caffe,TensorFlow,Mxnet,Pytorch 等主流深度学习框架,由英伟达开发维护。业界几乎所有 GPU 推理业务都在使用TensorRT。

但是 TensorRT 也存在缺陷,即它的部署流程比较繁琐,因此算法工程师提供的模型需要交由系统工程师来部署上线,非常耗时耗力。在传统的 TensorRT 工作流里,手写插件往往是最耗时的一部分。


TensorRT 手写算子插件难点
⦁    TensorRT 官方只支持很有限的常用算子(Conv/FC/BN/Relu…),对于不支持的算子,需要用户手写插件来实现;
⦁    插件的编写需要 GPU 和 cuda 知识,英伟达的工程师也通常需要 1~2 周时间来编写一个算子实现;模型中如果包含多个不支持算子,就需要更多时间来逐个编写和调试插件。

TPAT 概览

TPAT 实现了 TensorRT 插件的全自动生成,TensorRT 的部署和上线能基本流程化不再需要人工参与。手写插件的步骤将由 TPAT 代替,TPAT 全自动生成一个算子插件耗时仅需要 30-60 分钟的时间(该时间用于搜索算子的高性能 CUDA Kernel),TensorRT 会因此成为一个真正端到端的推理框架。

TPAT 亮点

⦁    覆盖度:支持 onnx/tensorflow/pyTorch 所有的算子
⦁    全自动:端到端全自动生成用户指定的 TensorRT Plugin
⦁    高性能:大部分算子上性能超越手写 Plugin

架构设计

TPAT 接受用户输入的 ONNX-Model,指定需要生成 TensorRT Plugin 的算子和 batch size,基于 TVM 深度学习编译器,对固定形状的算子进行 AutoTune,自动生成高性能的 CUDA Kernel. 将 CUDA Kernel 和 Runtime 必要的参数填充进 TensorRT Plugin 模板,生成动态链接库,可以直接加载到 TensorRT 运行。

TPAT 部分算子性能数据

使用 TPAT 自动生成 TensorRT-7.2 不支持的算子,并且用 TPAT 优化 TensorRT-7.2 原生实现性能较差的算子;

对比手写 Plugin

优化 TensorRT 原生算子

对内部业务模型里的部分算子进行了测试,TPAT 的性能几乎全面超越 CUDA 工程师手写,并且端到端的设计能够大幅减少人力投入;对于 TensorRT 原生的算子实现,TPAT 的表现也并不逊色,AutoTune 的特点能够优化 TensorRT 里表现不那么好的原生算子实现。

TPAT 开源

后续 TPAT 的开源计划:

  • 对于算子的多精度进行支持,包括 Float16,Int8.
  • 利用 TPAT 进行子图的优化
  • 对于动态形状的支持

附:TPAT 使用案例

使用 TPAT 支持 Onehot 算子(TensorRT-7.2.2.3)

  • 输入包括了 onehot 算子的 ONNX_Model、Onehot 算子的名字、batch_size
  • TPAT借助 TVM的Relay 和 AutoScheduler 组件,生成高性能的 CUDA Kernel;
  • 经过模板填充后直接生成可用的 onehot 算子 Plugin 的动态链接库。

展开阅读全文
点击加入讨论🔥(3) 发布并加入讨论🔥
3 评论
9 收藏
分享
返回顶部
顶部