TPAT是由腾讯联合英伟达开发的TensorRT插件自动生成工具。
TensorRT是当前应用最广的GPU推理框架，但由于支持的算子数量有限，用户面临手写插件以支持算子的痛点。而TPAT能够支持开放神经网络交换(ONNX)格式所有的算子，端到端生成TensorRT插件，在解放人力成本的同时，性能对比手写毫不逊色。
TPAT实现了TensorRT插件的全自动生成，TensorRT的部署和上线能基本流程化不再需要人工参与。手写插件的步骤将由TPAT代替，TPAT全自动生成一个算子插件耗时仅需要30-60分钟的时间（该时间用于搜索算子的高性能CUDAKernel）。
架构设计TPAT接受用户输入的ONNX-Model，指定需要生成TensorRTPlugin的算子和batchsize，基于TVM深度学习编译器，对固定形状的算子进行AutoTune，自动生成高性能的CUDAKernel.将CUDAKernel和Runtime必要的参数填充进TensorRTPlugin模板，生成动态链接库，可以直接加载到TensorRT运行。
TPAT部分算子性能数据使用TPAT自动生成TensorRT-7.2不支持的算子，并且用TPAT优化TensorRT-7.2原生实现性能较差的算子；对比手写Plugin优化TensorRT原生算子TPAT使用案例使用TPAT支持Onehot算子（TensorRT-7.2.2.3）输入包括了onehot算子的ONNX_Model、Onehot算子的名字、batch_sizeTPAT借助TVM的Relay和AutoScheduler组件，生成高性能的CUDAKernel；经过模板填充后直接生成可用的onehot算子Plugin的动态链接库。
评论